diff --git a/.github/actions/init-testing-instance-gpu/action.yml b/.github/actions/init-testing-instance-gpu/action.yml index 1e45fc2799..9d08b8089f 100644 --- a/.github/actions/init-testing-instance-gpu/action.yml +++ b/.github/actions/init-testing-instance-gpu/action.yml @@ -27,11 +27,11 @@ runs: - name: Install Rust deps shell: bash run: rustup install nightly-2025-02-14 --component rust-src - - name: Install OpenVM guest toolchain (nightly-2025-08-02) + - name: Install OpenVM guest toolchain (nightly-2026-01-18) shell: bash run: | - rustup toolchain install nightly-2025-08-02 - rustup component add rust-src --toolchain nightly-2025-08-02 + rustup toolchain install nightly-2026-01-18 + rustup component add rust-src --toolchain nightly-2026-01-18 - uses: taiki-e/install-action@nextest diff --git a/.github/actions/init-testing-instance/action.yml b/.github/actions/init-testing-instance/action.yml index c453d1b8ae..b617f36408 100644 --- a/.github/actions/init-testing-instance/action.yml +++ b/.github/actions/init-testing-instance/action.yml @@ -26,9 +26,9 @@ runs: - name: Install Rust deps shell: bash run: rustup install nightly-2025-02-14 --component rust-src - - name: Install OpenVM guest toolchain (nightly-2025-08-02) + - name: Install OpenVM guest toolchain (nightly-2026-01-18) shell: bash run: | - rustup toolchain install nightly-2025-08-02 - rustup component add rust-src --toolchain nightly-2025-08-02 + rustup toolchain install nightly-2026-01-18 + rustup component add rust-src --toolchain nightly-2026-01-18 - uses: taiki-e/install-action@nextest \ No newline at end of file diff --git a/.github/workflows/build-cache.yml b/.github/workflows/build-cache.yml index 1bcf9253d2..4839d69b52 100644 --- a/.github/workflows/build-cache.yml +++ b/.github/workflows/build-cache.yml @@ -26,18 +26,18 @@ jobs: run: rustup toolchain install nightly-2025-10-01 --component clippy,rustfmt - name: Install Rust toolchain run: rustup toolchain install nightly-2025-02-14 --component rust-src - - name: Install Rust toolchain 1.91 (stable) - run: rustup toolchain install 1.91 + - name: Install Rust toolchain 1.91.1 (stable, OpenVM v2 MSRV) + run: rustup toolchain install 1.91.1 - name: Set cargo to perform shallow clones run: echo "CARGO_NET_GIT_FETCH_WITH_CLI=true" >> $GITHUB_ENV - name: Format run: cargo fmt --all --check --verbose - - name: Cargo check with Rust 1.91 (default features) - run: cargo +1.91 check --all-targets + - name: Cargo check with Rust 1.91.1 (default features) + run: cargo +1.91.1 check --all-targets - name: Lint no default features run: cargo clippy --all --all-targets --no-default-features --profile pr-tests --verbose -- -D warnings - name: Build - run: cargo build --all-targets --features metrics --all --profile pr-tests --verbose + run: cargo build --all-targets --features metrics,aot --all --profile pr-tests --verbose ############################################################################### - name: Delete the old cache diff --git a/.github/workflows/nightly-tests.yml b/.github/workflows/nightly-tests.yml index e360099c3a..b4ebaa2f85 100644 --- a/.github/workflows/nightly-tests.yml +++ b/.github/workflows/nightly-tests.yml @@ -36,8 +36,8 @@ jobs: key: ${{ runner.os }}-cargo-pr-tests - name: Install Rust toolchain nightly-2025-10-01 (with clippy and rustfmt) run: rustup toolchain install nightly-2025-10-01 --component clippy,rustfmt,rust-src - - name: Install Rust toolchain 1.91 - run: rustup toolchain install 1.91 + - name: Install Rust toolchain 1.91.1 + run: rustup toolchain install 1.91.1 - name: Install riscv target run: rustup target add riscv32imac-unknown-none-elf --toolchain nightly-2025-10-01 - name: Install test dependencies @@ -78,10 +78,10 @@ jobs: run: cargo build --release -p powdr-openvm - name: Install cargo openvm - # Rust 1.91 is needed by fresher versions of dependencies of cargo-openvm. + # Rust 1.91.1 is needed by fresher versions of dependencies of cargo-openvm. run: | - rustup toolchain install 1.91 - cargo +1.91 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm + rustup toolchain install 1.91.1 + cargo +1.91.1 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm - name: Setup python venv run: | @@ -213,10 +213,10 @@ jobs: key: ${{ runner.os }}-cargo-release-apc-gpu-${{ hashFiles('**/Cargo.toml') }} - name: Install cargo openvm - # Rust 1.91 is needed by fresher versions of dependencies of cargo-openvm. + # Rust 1.91.1 is needed by fresher versions of dependencies of cargo-openvm. run: | - rustup toolchain install 1.91 - cargo +1.91 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm + rustup toolchain install 1.91.1 + cargo +1.91.1 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm - name: Setup python venv run: | diff --git a/.github/workflows/post-merge-tests.yml b/.github/workflows/post-merge-tests.yml index a03a89045b..3f97eba710 100644 --- a/.github/workflows/post-merge-tests.yml +++ b/.github/workflows/post-merge-tests.yml @@ -37,10 +37,10 @@ jobs: run: cargo build --release -p powdr-openvm - name: Install cargo openvm - # Rust 1.91 is needed by fresher versions of dependencies of cargo-openvm. + # Rust 1.91.1 is needed by fresher versions of dependencies of cargo-openvm. run: | - rustup toolchain install 1.91 - cargo +1.91 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm + rustup toolchain install 1.91.1 + cargo +1.91.1 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm - name: Run keccak with 100 APCs run: /usr/bin/time -v cargo run --bin powdr_openvm_riscv -r prove guest-keccak --input 10000 --autoprecompiles 100 --recursion diff --git a/.github/workflows/pr-tests-with-secrets.yml b/.github/workflows/pr-tests-with-secrets.yml index ab9e732906..237e21b289 100644 --- a/.github/workflows/pr-tests-with-secrets.yml +++ b/.github/workflows/pr-tests-with-secrets.yml @@ -50,10 +50,10 @@ jobs: run: cargo build --release -p powdr-openvm - name: Install cargo openvm - # Rust 1.91 is needed by fresher versions of dependencies of cargo-openvm. + # Rust 1.91.1 is needed by fresher versions of dependencies of cargo-openvm. run: | - rustup toolchain install 1.91 - cargo +1.91 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm + rustup toolchain install 1.91.1 + cargo +1.91.1 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm - name: Patch benchmark uses: ./.github/actions/patch-openvm-reth-benchmark @@ -106,10 +106,10 @@ jobs: run: cargo build --release -p powdr-openvm - name: Install cargo openvm - # Rust 1.91 is needed by fresher versions of dependencies of cargo-openvm. + # Rust 1.91.1 is needed by fresher versions of dependencies of cargo-openvm. run: | - rustup toolchain install 1.91 - cargo +1.91 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm + rustup toolchain install 1.91.1 + cargo +1.91.1 install --git 'http://github.com/powdr-labs/openvm.git' --rev "v1.4.2-powdr-rc.4" --locked cargo-openvm - name: Setup python venv run: | diff --git a/.github/workflows/pr-tests.yml b/.github/workflows/pr-tests.yml index ae9f9eab8c..ab9d25e305 100644 --- a/.github/workflows/pr-tests.yml +++ b/.github/workflows/pr-tests.yml @@ -47,18 +47,18 @@ jobs: ##### The block below is shared between cache build and PR build workflows ##### - name: Install Rust toolchain nightly-2025-10-01 (with clippy and rustfmt) run: rustup toolchain install nightly-2025-10-01 --component clippy,rustfmt - - name: Install Rust toolchain 1.91 (stable) - run: rustup toolchain install 1.91 + - name: Install Rust toolchain 1.91.1 (stable) + run: rustup toolchain install 1.91.1 - name: Set cargo to perform shallow clones run: echo "CARGO_NET_GIT_FETCH_WITH_CLI=true" >> $GITHUB_ENV - name: Format run: cargo fmt --all --check --verbose - - name: Cargo check with Rust 1.91 (default features) - run: cargo +1.91 check --all-targets + - name: Cargo check with Rust 1.91.1 (default features) + run: cargo +1.91.1 check --all-targets - name: Lint no default features - run: cargo clippy --all --all-targets --features metrics --profile pr-tests --verbose -- -D warnings + run: cargo clippy --all --all-targets --features metrics,aot --profile pr-tests --verbose -- -D warnings - name: Build (CPU) - run: cargo build --all-targets --features metrics --all --profile pr-tests --verbose + run: cargo build --all-targets --features metrics,aot --all --profile pr-tests --verbose ############################################################################### - uses: taiki-e/install-action@nextest @@ -161,14 +161,14 @@ jobs: - name: Install Rust toolchain nightly-2025-10-01 (with clippy and rustfmt) run: rustup toolchain install nightly-2025-10-01 --component clippy,rustfmt - - name: Install Rust toolchain 1.91 (stable) - run: rustup toolchain install 1.91 + - name: Install Rust toolchain 1.91.1 (stable) + run: rustup toolchain install 1.91.1 - name: Set cargo to perform shallow clones run: echo "CARGO_NET_GIT_FETCH_WITH_CLI=true" >> $GITHUB_ENV - name: Format run: cargo fmt --all --check --verbose - - name: Cargo check with Rust 1.91 (all features) - run: cargo +1.91 check --all-targets + - name: Cargo check with Rust 1.91.1 (all features) + run: cargo +1.91.1 check --all-targets - name: Lint no default features run: cargo clippy --all --all-targets --features cuda,metrics,aot --profile pr-tests --verbose -- -D warnings - name: Build (GPU) @@ -207,7 +207,10 @@ jobs: - name: Init testing instance (GPU) uses: ./.github/actions/init-testing-instance-gpu - name: Run quick GPU tests from powdr-openvm-riscv only - run: cargo nextest run --archive-file tests_gpu.tar.zst --workspace-remap . --verbose --no-tests=warn + # Cap parallelism so concurrent GPU allocations fit in ~20 GiB of device + # memory on our shared runner. Heavy APC tests can each consume several + # GiB, so we run at most 4 at a time. + run: cargo nextest run --archive-file tests_gpu.tar.zst --workspace-remap . --verbose --no-tests=warn --test-threads=4 # run: cargo nextest run --archive-file tests_gpu.tar.zst --workspace-remap . --verbose --partition count:"${{ matrix.test }}"/7 --no-tests=warn # NOTE: test_apc_reth_app_proof has been moved to pr-tests-with-secrets.yml diff --git a/Cargo.toml b/Cargo.toml index 9eb2950a12..590508ce0a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,54 +49,54 @@ powdr-openvm-riscv-hints-transpiler = { path = "./openvm-riscv/extensions/hints- powdr-openvm-riscv-hints-circuit = { path = "./openvm-riscv/extensions/hints-circuit", version = "0.1.4" } # openvm -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-build = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-rv32im-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-rv32im-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-rv32im-guest = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", default-features = false } -openvm-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-circuit-derive = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-circuit-primitives = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-circuit-primitives-derive = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-instructions = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-instructions-derive = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-sdk = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", default-features = false, features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-build = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-rv32im-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-rv32im-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-rv32im-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", default-features = false } +openvm-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-circuit-derive = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-circuit-primitives = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-circuit-primitives-derive = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-instructions = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-instructions-derive = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-sdk = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", default-features = false, features = [ "parallel", "jemalloc", - "nightly-features", - "evm-prove", ] } -openvm-ecc-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-ecc-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-keccak256-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-keccak256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-sha256-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-sha256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-algebra-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-algebra-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-bigint-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-bigint-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-pairing-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-pairing-transpiler = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-native-circuit = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", default-features = false } -openvm-native-recursion = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", default-features = false } -openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-custom-insn = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm-sdk-config = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-ecc-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-ecc-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-keccak256-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-keccak256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-sha2-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-sha2-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-algebra-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-algebra-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-bigint-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-bigint-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-pairing-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-pairing-transpiler = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-native-circuit = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", default-features = false } +openvm-native-recursion = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", default-features = false } +openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-custom-insn = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-poseidon2-air = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } # stark-backend -openvm-stark-sdk = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "v1.2.2-powdr-2026-03-20", default-features = false, features = [ +openvm-stark-sdk = { git = "https://github.com/powdr-labs/stark-backend.git", tag = "v2.0.0-beta.2-powdr", default-features = false, features = [ "parallel", "jemalloc", - "nightly-features", ] } -openvm-stark-backend = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "v1.2.2-powdr-2026-03-20", default-features = false, features = [ +openvm-stark-backend = { git = "https://github.com/powdr-labs/stark-backend.git", tag = "v2.0.0-beta.2-powdr", default-features = false, features = [ "parallel", "jemalloc", ] } -openvm-cuda-backend = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "v1.2.2-powdr-2026-03-20", default-features = false } -openvm-cuda-builder = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "v1.2.2-powdr-2026-03-20", default-features = false } -openvm-cuda-common = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "v1.2.2-powdr-2026-03-20", default-features = false } +openvm-cpu-backend = { git = "https://github.com/powdr-labs/stark-backend.git", tag = "v2.0.0-beta.2-powdr", default-features = false } +openvm-cuda-backend = { git = "https://github.com/powdr-labs/stark-backend.git", tag = "v2.0.0-beta.2-powdr", default-features = false } +openvm-cuda-builder = { git = "https://github.com/powdr-labs/stark-backend.git", tag = "v2.0.0-beta.2-powdr", default-features = false } +openvm-cuda-common = { git = "https://github.com/powdr-labs/stark-backend.git", tag = "v2.0.0-beta.2-powdr", default-features = false } # external dependencies num-traits = "0.2.19" @@ -131,7 +131,8 @@ debug-assertions = true overflow-checks = true panic = 'unwind' incremental = true # This is true because target is cached -codegen-units = 256 +codegen-units = 16 +lto = "thin" [profile.release-with-debug] inherits = "release" @@ -142,15 +143,17 @@ print_stdout = "deny" uninlined_format_args = "deny" iter_over_hash_type = "deny" -# Uncomment both patches below for local stark-backend and openvm. +# Local development patches - uncomment both sections together. # The local openvm also needs to have stark-backend patched so all types match. # [patch."https://github.com/powdr-labs/stark-backend.git"] # openvm-stark-sdk = { path = "../stark-backend/crates/stark-sdk", default-features = false } # openvm-stark-backend = { path = "../stark-backend/crates/stark-backend", default-features = false } +# openvm-codec-derive = { path = "../stark-backend/crates/stark-backend/codec-derive", default-features = false } # openvm-cuda-backend = { path = "../stark-backend/crates/cuda-backend", default-features = false } # openvm-cuda-builder = { path = "../stark-backend/crates/cuda-builder", default-features = false } # openvm-cuda-common = { path = "../stark-backend/crates/cuda-common", default-features = false } - +# openvm-cpu-backend = { path = "../stark-backend/crates/cpu-backend", default-features = false } +# # [patch."https://github.com/powdr-labs/openvm.git"] # openvm = { path = "../openvm/crates/toolchain/openvm" } # openvm-build = { path = "../openvm/crates/toolchain/build" } @@ -162,22 +165,25 @@ iter_over_hash_type = "deny" # openvm-circuit-derive = { path = "../openvm/crates/vm/derive" } # openvm-circuit-primitives = { path = "../openvm/crates/circuits/primitives" } # openvm-circuit-primitives-derive = { path = "../openvm/crates/circuits/primitives/derive" } -# openvm-instructions = { path = "../openvm/crates/toolchain/instructions" } -# openvm-instructions-derive = { path = "../openvm/crates/toolchain/instructions/derive" } # openvm-sdk = { path = "../openvm/crates/sdk" } +# openvm-sdk-config = { path = "../openvm/crates/sdk-config" } # openvm-ecc-circuit = { path = "../openvm/extensions/ecc/circuit" } # openvm-ecc-transpiler = { path = "../openvm/extensions/ecc/transpiler" } # openvm-keccak256-circuit = { path = "../openvm/extensions/keccak256/circuit" } # openvm-keccak256-transpiler = { path = "../openvm/extensions/keccak256/transpiler" } -# openvm-sha256-circuit = { path = "../openvm/extensions/sha256/circuit" } -# openvm-sha256-transpiler = { path = "../openvm/extensions/sha256/transpiler" } +# openvm-sha2-circuit = { path = "../openvm/extensions/sha2/circuit" } +# openvm-sha2-transpiler = { path = "../openvm/extensions/sha2/transpiler" } # openvm-algebra-circuit = { path = "../openvm/extensions/algebra/circuit" } # openvm-algebra-transpiler = { path = "../openvm/extensions/algebra/transpiler" } # openvm-bigint-circuit = { path = "../openvm/extensions/bigint/circuit" } # openvm-bigint-transpiler = { path = "../openvm/extensions/bigint/transpiler" } # openvm-pairing-circuit = { path = "../openvm/extensions/pairing/circuit" } # openvm-pairing-transpiler = { path = "../openvm/extensions/pairing/transpiler" } -# openvm-native-circuit = { path = "../openvm/extensions/native/circuit" } -# openvm-native-recursion = { path = "../openvm/extensions/native/recursion" } # openvm-platform = { path = "../openvm/crates/toolchain/platform" } # openvm-custom-insn = { path = "../openvm/crates/toolchain/custom_insn" } +# openvm-instructions = { path = "../openvm/crates/toolchain/instructions" } +# openvm-instructions-derive = { path = "../openvm/crates/toolchain/instructions/derive" } +# openvm-poseidon2-air = { path = "../openvm/crates/circuits/poseidon2-air" } +# continuations-v2 = { path = "../openvm/crates/continuations-v2" } +# recursion-circuit = { path = "../openvm/crates/recursion" } +# openvm-deferral-circuit = { path = "../openvm/extensions/deferral/circuit" } diff --git a/autoprecompiles/src/expression.rs b/autoprecompiles/src/expression.rs index 9aa13075ff..364c35a26e 100644 --- a/autoprecompiles/src/expression.rs +++ b/autoprecompiles/src/expression.rs @@ -176,20 +176,22 @@ where } } -/// Evaluates an `AlgebraicExpression` to a concrete value by subsituting the polynomial references by known values where known value is looked up via a column index mapping. +/// Evaluates an `AlgebraicExpression` to a concrete value by substituting the polynomial +/// references by known values where known value is looked up via a dense Vec index mapping. pub struct MappingRowEvaluator<'a, F> where F: Add + Sub + Mul + Neg + Copy, { pub row: &'a [F], - pub witness_id_to_index: &'a BTreeMap, + /// Dense Vec indexed by polynomial ID -> column index in row. + pub witness_id_to_index: &'a [usize], } impl<'a, F> MappingRowEvaluator<'a, F> where F: Add + Sub + Mul + Neg + Copy, { - pub fn new(row: &'a [F], witness_id_to_index: &'a BTreeMap) -> Self { + pub fn new(row: &'a [F], witness_id_to_index: &'a [usize]) -> Self { Self { row, witness_id_to_index, @@ -206,7 +208,7 @@ where } fn eval_var(&self, algebraic_var: &AlgebraicReference) -> F { - let index = self.witness_id_to_index[&(algebraic_var.id)]; + let index = self.witness_id_to_index[algebraic_var.id as usize]; self.row[index] } } @@ -250,3 +252,244 @@ where (*self.witness.get(&algebraic_var.id).unwrap()).into() } } + +use powdr_expression::{AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicUnaryOperator}; + +/// Pre-compiled expression for fast per-row evaluation without recursive AST walking. +/// +/// At build time, each `AlgebraicExpression` is analyzed once and compiled into the most +/// efficient variant. The specialized variants avoid the recursive enum-matching, closure +/// dispatch, and trait-object overhead of `AlgebraicExpression::to_expression()`. +pub enum CompiledExpr { + /// Expression is a compile-time constant. Zero cost per row. + Constant(F), + /// Expression is a single column read: `row[col_idx]`. One array access per row. + DirectLoad(usize), + /// Expression is `constant + sum(coeff * row[col_idx]) + sum(coeff * row[a] * row[b])`. + /// Covers both linear (degree-1) and quadratic (degree-2) expressions. + /// Linear expressions have an empty `products` vec. + Polynomial { + constant: F, + /// Linear terms: (col_idx, coefficient). + linear: Vec<(usize, F)>, + /// Quadratic terms: (col_idx_a, col_idx_b, coefficient). + products: Vec<(usize, usize, F)>, + }, +} + +/// Intermediate decomposition: constant + linear terms + quadratic products. +/// Used during compilation; converted to `CompiledExpr` afterwards. +struct Decomposition { + constant: F, + linear: Vec<(usize, F)>, + products: Vec<(usize, usize, F)>, +} + +impl CompiledExpr +where + F: Add + Sub + Mul + Neg + Copy + PartialEq, +{ + /// Compile an expression at build time into a fast-eval form. + /// `id_to_idx` is a dense Vec indexed directly by poly ID. + /// `zero` and `one` are the field's additive and multiplicative identities. + pub fn compile(expr: &AlgebraicExpression, id_to_idx: &[usize], zero: F, one: F) -> Self { + let d = Self::decompose(expr, id_to_idx, zero, one); + if d.linear.is_empty() && d.products.is_empty() { + CompiledExpr::Constant(d.constant) + } else if d.linear.len() == 1 + && d.products.is_empty() + && d.constant == zero + && d.linear[0].1 == one + { + CompiledExpr::DirectLoad(d.linear[0].0) + } else { + CompiledExpr::Polynomial { + constant: d.constant, + linear: d.linear, + products: d.products, + } + } + } + + /// Evaluate the compiled expression against a row slice. + #[inline(always)] + pub fn eval(&self, row: &[F]) -> F { + match self { + CompiledExpr::Constant(c) => *c, + CompiledExpr::DirectLoad(idx) => row[*idx], + CompiledExpr::Polynomial { + constant, + linear, + products, + } => { + let lin = linear + .iter() + .fold(*constant, |acc, &(idx, coeff)| acc + coeff * row[idx]); + products + .iter() + .fold(lin, |acc, &(a, b, coeff)| acc + coeff * row[a] * row[b]) + } + } + } + + /// Decompose an expression into constant + linear + quadratic components. + /// Panics if the expression is degree > 2. + fn decompose( + expr: &AlgebraicExpression, + id_to_idx: &[usize], + zero: F, + one: F, + ) -> Decomposition { + use powdr_expression::AlgebraicExpression as AE; + match expr { + AE::Number(c) => Decomposition { + constant: *c, + linear: vec![], + products: vec![], + }, + AE::Reference(r) => Decomposition { + constant: zero, + linear: vec![(id_to_idx[r.id as usize], one)], + products: vec![], + }, + AE::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) => { + let l = Self::decompose(left, id_to_idx, zero, one); + let r = Self::decompose(right, id_to_idx, zero, one); + match op { + AlgebraicBinaryOperator::Add => Decomposition { + constant: l.constant + r.constant, + linear: l.linear.into_iter().chain(r.linear).collect(), + products: l.products.into_iter().chain(r.products).collect(), + }, + AlgebraicBinaryOperator::Sub => Decomposition { + constant: l.constant - r.constant, + linear: l + .linear + .into_iter() + .chain(r.linear.into_iter().map(|(i, c)| (i, -c))) + .collect(), + products: l + .products + .into_iter() + .chain(r.products.into_iter().map(|(a, b, c)| (a, b, -c))) + .collect(), + }, + AlgebraicBinaryOperator::Mul => Self::mul_decompositions(l, r, zero), + } + } + AE::UnaryOperation(powdr_expression::AlgebraicUnaryOperation { op, expr }) => { + match op { + AlgebraicUnaryOperator::Minus => { + let d = Self::decompose(expr, id_to_idx, zero, one); + Decomposition { + constant: -d.constant, + linear: d.linear.into_iter().map(|(i, c)| (i, -c)).collect(), + products: d.products.into_iter().map(|(a, b, c)| (a, b, -c)).collect(), + } + } + } + } + } + } + + /// Multiply two decompositions. Supports constant * anything and + /// linear * linear (producing quadratic products). + /// Panics if the result would be degree > 2. + fn mul_decompositions(l: Decomposition, r: Decomposition, zero: F) -> Decomposition { + // If either side is purely constant, scale the other. + if l.linear.is_empty() && l.products.is_empty() { + return Decomposition { + constant: l.constant * r.constant, + linear: r + .linear + .into_iter() + .map(|(i, c)| (i, c * l.constant)) + .collect(), + products: r + .products + .into_iter() + .map(|(a, b, c)| (a, b, c * l.constant)) + .collect(), + }; + } + if r.linear.is_empty() && r.products.is_empty() { + return Decomposition { + constant: l.constant * r.constant, + linear: l + .linear + .into_iter() + .map(|(i, c)| (i, c * r.constant)) + .collect(), + products: l + .products + .into_iter() + .map(|(a, b, c)| (a, b, c * r.constant)) + .collect(), + }; + } + // Both sides have variable terms. Products of quadratic with anything → degree > 2. + assert!( + l.products.is_empty() && r.products.is_empty(), + "Bus interaction expression is degree > 2. This is unexpected." + ); + // (c_l + sum(a_i * x_i)) * (c_r + sum(b_j * y_j)) + // = c_l*c_r + c_l*sum(b_j*y_j) + c_r*sum(a_i*x_i) + sum(a_i*b_j * x_i*y_j) + let mut linear = Vec::new(); + let mut products = Vec::new(); + + // c_l * linear_r + if l.constant != zero { + linear.extend(r.linear.iter().map(|&(i, c)| (i, c * l.constant))); + } + // c_r * linear_l + if r.constant != zero { + linear.extend(l.linear.iter().map(|&(i, c)| (i, c * r.constant))); + } + // linear_l * linear_r → quadratic products + for &(i, ci) in &l.linear { + for &(j, cj) in &r.linear { + products.push((i, j, ci * cj)); + } + } + + Decomposition { + constant: l.constant * r.constant, + linear, + products, + } + } +} + +/// Pre-compiled bus interaction for fast per-row evaluation. +pub struct CompiledBusInteraction { + pub id: u64, + pub mult: CompiledExpr, + pub args: Vec>, +} + +impl CompiledBusInteraction +where + F: Add + Sub + Mul + Neg + Copy + PartialEq, +{ + /// Compile all bus interactions from symbolic form. + /// `zero` and `one` are the field's additive and multiplicative identities. + pub fn compile_all( + interactions: &[SymbolicBusInteraction], + id_to_idx: &[usize], + zero: F, + one: F, + ) -> Vec { + interactions + .iter() + .map(|bi| CompiledBusInteraction { + id: bi.id, + mult: CompiledExpr::compile(&bi.mult, id_to_idx, zero, one), + args: bi + .args + .iter() + .map(|a| CompiledExpr::compile(a, id_to_idx, zero, one)) + .collect(), + }) + .collect() + } +} diff --git a/autoprecompiles/tests/optimizer.rs b/autoprecompiles/tests/optimizer.rs index 1dd3155037..ce1ff0bd57 100644 --- a/autoprecompiles/tests/optimizer.rs +++ b/autoprecompiles/tests/optimizer.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::fs; use std::path::Path; @@ -7,7 +8,7 @@ use powdr_autoprecompiles::bus_map::BusMap; use powdr_autoprecompiles::export::{ApcWithBusMap, SimpleInstruction}; use powdr_autoprecompiles::optimizer::optimize; use powdr_autoprecompiles::symbolic_machine::SymbolicMachine; -use powdr_autoprecompiles::{Apc, ColumnAllocator, DegreeBound}; +use powdr_autoprecompiles::{Apc, ColumnAllocator, DegreeBound, Substitution}; use powdr_number::BabyBearField; use powdr_openvm_bus_interaction_handler::memory_bus_interaction::OpenVmMemoryBusInteraction; use powdr_openvm_bus_interaction_handler::{ @@ -320,3 +321,592 @@ fn wasm_register_reuse() { "#]] .assert_debug_eq(&machine.constraints.len()); } + +/// Reusable JIT analysis function +#[allow(clippy::print_stdout)] +fn jit_analysis(apc_data: &ApcWithBusMap>, label: &str) { + let bus_map = &apc_data.bus_map; + + let pre_opt_machine: SymbolicMachine = apc_data.apc.machine.clone(); + let pre_opt_subs: Vec> = apc_data.apc.subs.clone(); + let pre_opt_cols: Vec<_> = pre_opt_machine.main_columns().sorted().collect(); + + println!("\n{}", "=".repeat(60)); + println!(" JIT Analysis: {label}"); + println!("{}", "=".repeat(60)); + println!("Pre-opt: {} cols, {} constraints, {} bus, {} instructions", + pre_opt_cols.len(), pre_opt_machine.constraints.len(), + pre_opt_machine.bus_interactions.len(), pre_opt_subs.len()); + + let machine: SymbolicMachine = apc_data.apc.machine.clone(); + let column_allocator = ColumnAllocator::from_max_poly_id_of_machine(&machine); + let (opt_machine, _) = optimize::<_, _, _, OpenVmMemoryBusInteraction<_, _>>( + machine, + OpenVmBusInteractionHandler::default(), + DEFAULT_DEGREE_BOUND, + bus_map, + column_allocator, + &mut Default::default(), + ) + .unwrap(); + + let opt_col_ids: std::collections::BTreeSet = + opt_machine.main_columns().map(|r| r.id).collect(); + let opt_cols: Vec<_> = opt_machine.main_columns().sorted().collect(); + + println!("Post-opt: {} cols, {} constraints, {} bus, {} derived", + opt_cols.len(), opt_machine.constraints.len(), + opt_machine.bus_interactions.len(), opt_machine.derived_columns.len()); + + let mut total_surviving_subs = 0; + let unused_instructions = pre_opt_subs + .iter() + .filter(|subs| subs.iter().all(|s| !opt_col_ids.contains(&s.apc_poly_id))) + .count(); + for subs in &pre_opt_subs { + total_surviving_subs += subs.iter().filter(|s| opt_col_ids.contains(&s.apc_poly_id)).count(); + } + println!("Column reduction: {} -> {} ({:.1}x)", + pre_opt_cols.len(), opt_cols.len(), + pre_opt_cols.len() as f64 / opt_cols.len().max(1) as f64); + println!("Surviving subs: {total_surviving_subs}/{}", pre_opt_cols.len()); + println!("Zero-sub instructions: {unused_instructions}/{}", pre_opt_subs.len()); + println!("Derived columns: {}", opt_machine.derived_columns.len()); + + // Width distribution + let width_distribution: BTreeMap = pre_opt_subs + .iter() + .map(|subs| subs.len()) + .fold(BTreeMap::new(), |mut acc, w| { *acc.entry(w).or_default() += 1; acc }); + println!("\nAIR types (by width): {:?}", width_distribution); + + // Build id-to-name map + let id_to_name: BTreeMap = pre_opt_cols + .iter() + .map(|c| (c.id, c.name.to_string())) + .collect(); + + // Global pattern classification + let mut global_patterns: BTreeMap<&str, usize> = BTreeMap::new(); + for subs in &pre_opt_subs { + for sub in subs { + if !opt_col_ids.contains(&sub.apc_poly_id) { continue; } + let col_name = id_to_name.get(&sub.apc_poly_id).cloned().unwrap_or_default(); + let base_name = col_name.rsplitn(2, '_').last().unwrap_or(&col_name); + let category = if base_name.contains("timestamp_lt_aux__lower_decomp") { + "COMPUTED: timestamp_decomp" + } else if base_name.contains("mem_ptr_limbs") { + "COMPUTED: mem_ptr_limbs" + } else if base_name.contains("write_data__") { + "COMPUTED: write_data" + } else if base_name.contains("flags__") || base_name.contains("is_valid") + || base_name.contains("is_load") || base_name.contains("needs_write") + || base_name.contains("cmp_result") + { + "COMPUTED: flags/opcode_dep" + } else if base_name.contains("sign_xor") || base_name.contains("_inv") + || base_name.contains("lt_marker") || base_name.contains("lt_diff") + || base_name.contains("r_prime") || base_name.contains("zero_divisor") + || base_name.contains("r_zero") || base_name.contains("q_sign") + || base_name.contains("c_sign") || base_name.contains("c_sum_inv") + || base_name.contains("r_sum_inv") + { + "COMPUTED: division/comparison aux" + } else if base_name.contains("shift_amount") || base_name.contains("bit_shift") + || base_name.contains("_shifted") || base_name.contains("limb_shift") + { + "COMPUTED: shift aux" + } else { + "DIRECT: copy/byte_decomp" + }; + *global_patterns.entry(category).or_default() += 1; + } + } + let total: usize = global_patterns.values().sum(); + println!("\nColumn computation categories:"); + let mut sorted: Vec<_> = global_patterns.iter().collect(); + sorted.sort_by(|a, b| b.1.cmp(a.1)); + for (cat, count) in &sorted { + println!(" {count:>5} ({:>5.1}%) {cat}", **count as f64 / total as f64 * 100.0); + } + println!(" Total: {total}"); + + // Derived column expressions + if !opt_machine.derived_columns.is_empty() { + println!("\nDerived columns:"); + for (i, d) in opt_machine.derived_columns.iter().enumerate().take(5) { + println!(" {i}: {} = {}", d.variable, d.computation_method); + } + if opt_machine.derived_columns.len() > 5 { + println!(" ... and {} more", opt_machine.derived_columns.len() - 5); + } + } + + // Constraint-only columns + let substituted_col_ids: std::collections::BTreeSet = pre_opt_subs + .iter() + .flat_map(|subs| subs.iter()) + .filter(|s| opt_col_ids.contains(&s.apc_poly_id)) + .map(|s| s.apc_poly_id) + .collect(); + let derived_col_ids: std::collections::BTreeSet = opt_machine + .derived_columns.iter().map(|d| d.variable.id).collect(); + let constraint_only: usize = opt_col_ids.iter() + .filter(|id| !substituted_col_ids.contains(id) && !derived_col_ids.contains(id)) + .count(); + if constraint_only > 0 { + println!("\nWARNING: {constraint_only} constraint-only columns (not substituted, not derived)!"); + for id in opt_col_ids.iter() + .filter(|id| !substituted_col_ids.contains(id) && !derived_col_ids.contains(id)) + .take(10) + { + let name = opt_cols.iter().find(|c| c.id == *id).map(|c| c.name.as_str()).unwrap_or("?"); + println!(" {} (id={})", name, id); + } + } +} + +/// Detailed analysis of a keccak APC for JIT trace generation feasibility. +#[test] +#[allow(clippy::print_stdout)] +fn jit_tracegen_analysis_keccak() { + let apc_data = import_apc_from_gzipped_json("tests/keccak_apc_pre_opt.json.gz"); + jit_analysis(&apc_data, "keccak"); +} + +/// Same analysis for ecrecover (different program type) +#[test] +#[allow(clippy::print_stdout)] +fn jit_tracegen_analysis_ecrecover() { + let apc_data = import_apc_from_gzipped_json("tests/ecrecover_apc_pre_opt.json.gz"); + jit_analysis(&apc_data, "ecrecover"); +} + +/// Same analysis for sha256 +#[test] +#[allow(clippy::print_stdout)] +fn jit_tracegen_analysis_sha256() { + let apc_data = import_apc_from_gzipped_json("tests/sha256_apc_pre_opt.json.gz"); + jit_analysis(&apc_data, "sha256"); +} + +/// Same analysis for single div_nondet +#[test] +#[allow(clippy::print_stdout)] +fn jit_tracegen_analysis_div() { + let apc_data = import_apc_from_gzipped_json("tests/single_div_nondet.json.gz"); + jit_analysis(&apc_data, "div_nondet"); +} + +/// Same analysis for reth_op +#[test] +#[allow(clippy::print_stdout)] +fn jit_tracegen_analysis_reth() { + let apc_data = import_apc_from_gzipped_json("tests/apc_reth_op_bug.json.gz"); + jit_analysis(&apc_data, "reth_op"); +} + +/// Detailed analysis of a keccak APC for JIT trace generation feasibility. +/// Dumps: +/// - Pre/post-optimization column counts and names +/// - Substitution mapping (which original columns survive) +/// - Derived column expressions (computed columns) +/// - Per-instruction substitution breakdown +/// - Column classification (substituted vs derived vs constrained-only) +#[test] +#[allow(clippy::print_stdout)] +fn jit_tracegen_analysis_keccak_detail() { + let apc_data = import_apc_from_gzipped_json("tests/keccak_apc_pre_opt.json.gz"); + let bus_map = &apc_data.bus_map; + + let pre_opt_machine: SymbolicMachine = apc_data.apc.machine.clone(); + let pre_opt_subs: Vec> = apc_data.apc.subs.clone(); + + // Pre-optimization stats + let pre_opt_cols: Vec<_> = pre_opt_machine.main_columns().sorted().collect(); + println!("=== PRE-OPTIMIZATION ==="); + println!("Total columns: {}", pre_opt_cols.len()); + println!("Total constraints: {}", pre_opt_machine.constraints.len()); + println!( + "Total bus interactions: {}", + pre_opt_machine.bus_interactions.len() + ); + println!( + "Total derived columns: {}", + pre_opt_machine.derived_columns.len() + ); + println!( + "Instructions in block (subs count): {}", + pre_opt_subs.len() + ); + + // Per-instruction substitution counts + println!("\n=== PRE-OPT SUBSTITUTIONS (per instruction) ==="); + for (i, subs) in pre_opt_subs.iter().enumerate() { + println!( + " Instruction {i}: {} substitutions (original col indices: [{}])", + subs.len(), + subs.iter() + .map(|s| format!( + "orig_idx={} -> apc_id={}", + s.original_poly_index, s.apc_poly_id + )) + .take(5) // limit output + .join(", ") + ); + if subs.len() > 5 { + println!(" ... and {} more", subs.len() - 5); + } + } + + // Now optimize + let machine: SymbolicMachine = apc_data.apc.machine; + let column_allocator = ColumnAllocator::from_max_poly_id_of_machine(&machine); + let (opt_machine, _opt_col_alloc) = optimize::<_, _, _, OpenVmMemoryBusInteraction<_, _>>( + machine, + OpenVmBusInteractionHandler::default(), + DEFAULT_DEGREE_BOUND, + bus_map, + column_allocator, + &mut Default::default(), + ) + .unwrap(); + + // Build the Apc to get post-opt substitutions + // Note: Apc::new filters subs to only keep those referenced in the optimized machine + // We can't call Apc::new without the block, so instead we manually check which pre-opt + // subs survive by checking if their apc_poly_id is in the optimized machine's columns. + let opt_col_ids: std::collections::BTreeSet = + opt_machine.main_columns().map(|r| r.id).collect(); + + println!("\n=== POST-OPTIMIZATION ==="); + let opt_cols: Vec<_> = opt_machine.main_columns().sorted().collect(); + println!("Total columns: {}", opt_cols.len()); + println!("Total constraints: {}", opt_machine.constraints.len()); + println!( + "Total bus interactions: {}", + opt_machine.bus_interactions.len() + ); + println!( + "Total derived columns: {}", + opt_machine.derived_columns.len() + ); + + // Count surviving substitutions per instruction + println!("\n=== POST-OPT SURVIVING SUBSTITUTIONS (per instruction) ==="); + let mut total_surviving_subs = 0; + let mut total_original_subs = 0; + for (i, subs) in pre_opt_subs.iter().enumerate() { + let surviving: Vec<_> = subs + .iter() + .filter(|s| opt_col_ids.contains(&s.apc_poly_id)) + .collect(); + total_surviving_subs += surviving.len(); + total_original_subs += subs.len(); + if !surviving.is_empty() { + println!( + " Instruction {i}: {}/{} substitutions survive", + surviving.len(), + subs.len() + ); + } + } + println!( + "\n TOTAL: {total_surviving_subs}/{total_original_subs} substitutions survive optimization" + ); + + // Classify columns + let substituted_col_ids: std::collections::BTreeSet = pre_opt_subs + .iter() + .flat_map(|subs| subs.iter()) + .filter(|s| opt_col_ids.contains(&s.apc_poly_id)) + .map(|s| s.apc_poly_id) + .collect(); + + let derived_col_ids: std::collections::BTreeSet = opt_machine + .derived_columns + .iter() + .map(|d| d.variable.id) + .collect(); + + let constraint_only_col_ids: std::collections::BTreeSet = opt_col_ids + .iter() + .filter(|id| !substituted_col_ids.contains(id) && !derived_col_ids.contains(id)) + .cloned() + .collect(); + + println!("\n=== COLUMN CLASSIFICATION ==="); + println!( + "Substituted (copied from original traces): {}", + substituted_col_ids.len() + ); + println!( + "Derived (computed from other columns): {}", + derived_col_ids.len() + ); + println!( + "Constraint-only (not substituted, not derived): {}", + constraint_only_col_ids.len() + ); + + // Show derived column expressions (first 20) + println!("\n=== DERIVED COLUMN EXPRESSIONS (first 20) ==="); + for (i, d) in opt_machine.derived_columns.iter().take(20).enumerate() { + println!(" {i}: {} = {}", d.variable, d.computation_method); + } + if opt_machine.derived_columns.len() > 20 { + println!( + " ... and {} more", + opt_machine.derived_columns.len() - 20 + ); + } + + // Show constraint-only columns (these are the problematic ones for JIT - + // they aren't directly substituted and aren't derived, so their values come + // from somewhere else, likely the constraint solver) + println!("\n=== CONSTRAINT-ONLY COLUMNS (first 30) ==="); + let constraint_only_names: Vec<_> = opt_cols + .iter() + .filter(|c| constraint_only_col_ids.contains(&c.id)) + .take(30) + .collect(); + for col in &constraint_only_names { + println!(" {} (id={})", col.name, col.id); + } + if constraint_only_col_ids.len() > 30 { + println!( + " ... and {} more", + constraint_only_col_ids.len() - 30 + ); + } + + // Show the first 10 substituted column names for context + println!("\n=== SAMPLE SUBSTITUTED COLUMN NAMES (first 30) ==="); + let sub_col_names: Vec<_> = opt_cols + .iter() + .filter(|c| substituted_col_ids.contains(&c.id)) + .take(30) + .collect(); + for col in &sub_col_names { + println!(" {} (id={})", col.name, col.id); + } + + // Show first 10 constraints + println!("\n=== ALGEBRAIC CONSTRAINTS (first 10) ==="); + for (i, c) in opt_machine.constraints.iter().take(10).enumerate() { + println!(" {i}: {} = 0", c); + } + if opt_machine.constraints.len() > 10 { + println!( + " ... and {} more constraints", + opt_machine.constraints.len() - 10 + ); + } + + // Show first 10 bus interactions + println!("\n=== BUS INTERACTIONS (first 10) ==="); + for (i, bi) in opt_machine.bus_interactions.iter().take(10).enumerate() { + let bus_type = bus_map.bus_type(bi.id); + println!( + " {i}: bus={} ({}), mult={}, args=[{}]", + bi.id, + bus_type, + bi.mult, + bi.args.iter().join(", ") + ); + } + if opt_machine.bus_interactions.len() > 10 { + println!( + " ... and {} more bus interactions", + opt_machine.bus_interactions.len() - 10 + ); + } + + // JIT feasibility summary + println!("\n=== JIT FEASIBILITY SUMMARY ==="); + println!( + "Column reduction: {} -> {} ({:.1}x)", + pre_opt_cols.len(), + opt_cols.len(), + pre_opt_cols.len() as f64 / opt_cols.len() as f64 + ); + println!( + "Substituted columns (direct copy, trivial for JIT): {}", + substituted_col_ids.len() + ); + println!( + "Derived columns (algebraic formula, easy for JIT): {}", + derived_col_ids.len() + ); + println!( + "Constraint-only columns (need solver/witness gen, hard for JIT): {}", + constraint_only_col_ids.len() + ); + let trivial_pct = + (substituted_col_ids.len() + derived_col_ids.len()) as f64 / opt_cols.len() as f64 * 100.0; + println!( + "JIT-trivial columns: {:.1}% ({}/{})", + trivial_pct, + substituted_col_ids.len() + derived_col_ids.len(), + opt_cols.len() + ); + + // Check how many original instructions contribute no surviving subs + let unused_instructions = pre_opt_subs + .iter() + .filter(|subs| subs.iter().all(|s| !opt_col_ids.contains(&s.apc_poly_id))) + .count(); + println!( + "Original instructions with NO surviving substitutions: {}/{}", + unused_instructions, + pre_opt_subs.len() + ); + println!( + " -> These instructions' traces are generated but entirely discarded!" + ); + + // Count distinct instruction widths (proxy for AIR types) + let width_distribution: BTreeMap = pre_opt_subs + .iter() + .map(|subs| subs.len()) + .fold(BTreeMap::new(), |mut acc, w| { + *acc.entry(w).or_default() += 1; + acc + }); + println!("\n=== INSTRUCTION WIDTH DISTRIBUTION (proxy for AIR types) ==="); + for (width, count) in &width_distribution { + println!(" Width {width} columns: {count} instructions"); + } + println!(" Distinct widths: {}", width_distribution.len()); + + // Build a map from apc_poly_id to column name (from pre-opt machine) + let id_to_name: BTreeMap = pre_opt_cols + .iter() + .map(|c| (c.id, c.name.to_string())) + .collect(); + + // For each AIR width (proxy for AIR type), categorize surviving column names by pattern + println!("\n=== SURVIVING COLUMN PATTERNS BY AIR TYPE ==="); + for (air_width, _) in &width_distribution { + let mut pattern_counts: BTreeMap = BTreeMap::new(); + let mut sample_cols: BTreeMap> = BTreeMap::new(); + + for (i, subs) in pre_opt_subs.iter().enumerate() { + if subs.len() != *air_width { + continue; + } + for sub in subs { + if !opt_col_ids.contains(&sub.apc_poly_id) { + continue; + } + let col_name = id_to_name + .get(&sub.apc_poly_id) + .cloned() + .unwrap_or_else(|| format!("unknown_{}", sub.apc_poly_id)); + // Strip the instruction suffix (e.g., "_0", "_123") to get the base column name + let base_name = col_name.rsplitn(2, '_').last().unwrap_or(&col_name); + // Classify by pattern + let pattern = if base_name.contains("timestamp_lt_aux__lower_decomp") { + "timestamp_decomp".to_string() + } else if base_name.contains("prev_timestamp") { + "prev_timestamp (direct copy)".to_string() + } else if base_name.contains("from_state__timestamp") { + "from_state_timestamp (direct copy)".to_string() + } else if base_name.contains("from_state__pc") { + "from_state_pc (direct copy)".to_string() + } else if base_name.contains("mem_ptr_limbs") { + "mem_ptr_limbs (computed)".to_string() + } else if base_name.contains("read_data__") || base_name.contains("prev_data__") { + "memory_data (direct copy)".to_string() + } else if base_name.contains("rs1_data__") || base_name.contains("rs2_data__") { + "register_data (byte decomp)".to_string() + } else if base_name.contains("write_data__") { + "write_data (computed)".to_string() + } else if base_name.contains("flags__") || base_name.contains("is_valid") + || base_name.contains("is_load") + { + "flags/opcode_dep (computed)".to_string() + } else if base_name.contains("_ptr") { + "register_ptr (direct copy)".to_string() + } else if base_name.contains("imm") { + "immediate (direct copy)".to_string() + } else if base_name.contains("mem_as") { + "mem_as (direct copy)".to_string() + } else if base_name.contains("needs_write") { + "needs_write (computed)".to_string() + } else if base_name.contains("a__") || base_name.contains("b__") + || base_name.contains("c__") || base_name.contains("d__") + { + "alu_result_limbs (direct/computed)".to_string() + } else if base_name.contains("cmp_result") { + "cmp_result (computed)".to_string() + } else if base_name.contains("xor_result") + || base_name.contains("and_result") + || base_name.contains("or_result") + { + "bitwise_result (direct copy)".to_string() + } else { + format!("OTHER: {}", base_name) + }; + *pattern_counts.entry(pattern.clone()).or_default() += 1; + let samples = sample_cols.entry(pattern).or_default(); + if samples.len() < 3 { + samples.push((sub.original_poly_index, col_name.clone())); + } + } + } + println!("\n AIR width={air_width}:"); + let mut sorted: Vec<_> = pattern_counts.iter().collect(); + sorted.sort_by(|a, b| b.1.cmp(a.1)); + for (pattern, count) in &sorted { + let samples = sample_cols.get(*pattern).unwrap(); + let sample_str = samples + .iter() + .map(|(idx, name)| format!("col{}={}", idx, name)) + .collect::>() + .join(", "); + println!(" {count:>4}x {pattern} [{sample_str}]"); + } + } + + // Summary: total across all AIR types + println!("\n=== GLOBAL PATTERN SUMMARY ==="); + let mut global_patterns: BTreeMap = BTreeMap::new(); + for (i, subs) in pre_opt_subs.iter().enumerate() { + for sub in subs { + if !opt_col_ids.contains(&sub.apc_poly_id) { + continue; + } + let col_name = id_to_name + .get(&sub.apc_poly_id) + .cloned() + .unwrap_or_default(); + let base_name = col_name.rsplitn(2, '_').last().unwrap_or(&col_name); + let category = if base_name.contains("timestamp_lt_aux__lower_decomp") { + "COMPUTED: timestamp_decomp" + } else if base_name.contains("mem_ptr_limbs") { + "COMPUTED: mem_ptr_limbs" + } else if base_name.contains("write_data__") { + "COMPUTED: write_data" + } else if base_name.contains("flags__") + || base_name.contains("is_valid") + || base_name.contains("is_load") + || base_name.contains("needs_write") + || base_name.contains("cmp_result") + { + "COMPUTED: flags/opcode_dep" + } else { + "DIRECT: copy/byte_decomp" + }; + *global_patterns.entry(category.to_string()).or_default() += 1; + } + } + let total_surviving: usize = global_patterns.values().sum(); + for (cat, count) in &global_patterns { + println!( + " {count:>5} ({:>5.1}%) {cat}", + *count as f64 / total_surviving as f64 * 100.0 + ); + } + println!(" Total: {total_surviving}"); +} diff --git a/cli-openvm-riscv/Cargo.toml b/cli-openvm-riscv/Cargo.toml index 395909f6ff..c7f97858d0 100644 --- a/cli-openvm-riscv/Cargo.toml +++ b/cli-openvm-riscv/Cargo.toml @@ -7,8 +7,10 @@ homepage.workspace = true repository.workspace = true [features] -default = ["metrics"] +default = ["metrics", "cuda"] metrics = ["powdr-openvm/metrics", "openvm-sdk/metrics", "openvm-stark-backend/metrics", "openvm-stark-sdk/metrics"] +cuda = ["powdr-openvm-riscv/cuda"] +aot = ["powdr-openvm-riscv/aot"] [[bin]] name = "powdr_openvm_riscv" diff --git a/openvm-riscv/Cargo.toml b/openvm-riscv/Cargo.toml index 369e5f7210..8e680ca8a2 100644 --- a/openvm-riscv/Cargo.toml +++ b/openvm-riscv/Cargo.toml @@ -21,26 +21,26 @@ openvm-rv32im-transpiler.workspace = true openvm-rv32im-guest.workspace = true openvm-transpiler.workspace = true openvm-circuit.workspace = true +openvm-cpu-backend.workspace = true openvm-circuit-derive.workspace = true openvm-circuit-primitives.workspace = true openvm-circuit-primitives-derive.workspace = true openvm-instructions.workspace = true openvm-instructions-derive.workspace = true openvm-sdk.workspace = true +openvm-sdk-config.workspace = true openvm-ecc-circuit.workspace = true openvm-ecc-transpiler.workspace = true openvm-keccak256-circuit.workspace = true openvm-keccak256-transpiler.workspace = true -openvm-sha256-circuit.workspace = true -openvm-sha256-transpiler.workspace = true +openvm-sha2-circuit.workspace = true +openvm-sha2-transpiler.workspace = true openvm-algebra-circuit.workspace = true openvm-algebra-transpiler.workspace = true openvm-bigint-circuit.workspace = true openvm-bigint-transpiler.workspace = true openvm-pairing-circuit.workspace = true openvm-pairing-transpiler.workspace = true -openvm-native-circuit.workspace = true -openvm-native-recursion.workspace = true openvm-stark-sdk.workspace = true openvm-stark-backend.workspace = true diff --git a/openvm-riscv/extensions/hints-circuit/Cargo.toml b/openvm-riscv/extensions/hints-circuit/Cargo.toml index a6a6924f18..5b5bd054ee 100644 --- a/openvm-riscv/extensions/hints-circuit/Cargo.toml +++ b/openvm-riscv/extensions/hints-circuit/Cargo.toml @@ -21,5 +21,5 @@ powdr-openvm-riscv-hints-transpiler = { workspace = true } eyre.workspace = true crypto-bigint = "0.6.1" elliptic-curve = "0.13.8" -rand = { version = "0.8.5", default-features = false } +rand = { version = "0.9.2", default-features = false } serde.workspace = true diff --git a/openvm-riscv/extensions/hints-circuit/src/executors.rs b/openvm-riscv/extensions/hints-circuit/src/executors.rs index 7f351ba21f..7dfcbc9384 100644 --- a/openvm-riscv/extensions/hints-circuit/src/executors.rs +++ b/openvm-riscv/extensions/hints-circuit/src/executors.rs @@ -29,11 +29,7 @@ impl PhantomSubExecutor for ReverseBytesSubEx { // read memory let bytes = unsafe { memory.read::(RV32_MEMORY_AS, rs1) }; // write hint as bytes in reverse - let hint_bytes = bytes - .into_iter() - .rev() - .map(|b| F::from_canonical_u8(b)) - .collect(); + let hint_bytes = bytes.into_iter().rev().map(|b| F::from_u8(b)).collect(); streams.hint_stream = hint_bytes; Ok(()) } @@ -85,7 +81,7 @@ impl PhantomSubExecutor for K256InverseFieldSubEx { let inv_bytes = n_inv .to_be_bytes() .into_iter() - .map(|b| F::from_canonical_u8(b)) + .map(|b| F::from_u8(b)) .collect(); streams.hint_stream = inv_bytes; @@ -132,10 +128,7 @@ impl PhantomSubExecutor for K256InverseField10x26SubEx { let inv = elem.invert().normalize(); // okay to transmute in the opposite direction let inv_bytes: [u8; FIELD10X26_BYTES] = unsafe { std::mem::transmute(inv.0) }; - streams.hint_stream = inv_bytes - .into_iter() - .map(|b| F::from_canonical_u8(b)) - .collect(); + streams.hint_stream = inv_bytes.into_iter().map(|b| F::from_u8(b)).collect(); Ok(()) } @@ -190,7 +183,7 @@ impl PhantomSubExecutor for K256SqrtField10x26SubEx { .to_le_bytes() // indicates that a square root exists .into_iter() .chain(bytes) - .map(|b| F::from_canonical_u8(b)) + .map(|b| F::from_u8(b)) .collect(); } else { // Number is not square. @@ -204,7 +197,7 @@ impl PhantomSubExecutor for K256SqrtField10x26SubEx { .to_le_bytes() // indicate number is not square .into_iter() .chain(bytes) - .map(|b| F::from_canonical_u8(b)) + .map(|b| F::from_u8(b)) .collect(); } diff --git a/openvm-riscv/extensions/hints-circuit/src/lib.rs b/openvm-riscv/extensions/hints-circuit/src/lib.rs index b10c60006a..86a7d113c6 100644 --- a/openvm-riscv/extensions/hints-circuit/src/lib.rs +++ b/openvm-riscv/extensions/hints-circuit/src/lib.rs @@ -12,9 +12,8 @@ use openvm_circuit::derive::{ }; use openvm_circuit::system::phantom::PhantomExecutor; use openvm_instructions::PhantomDiscriminant; -use openvm_stark_backend::config::{StarkGenericConfig, Val}; use openvm_stark_backend::p3_field::{Field, PrimeField32}; -use openvm_stark_sdk::engine::StarkEngine; +use openvm_stark_backend::{StarkEngine, StarkProtocolConfig, Val}; use powdr_openvm_riscv_hints_transpiler::HintsPhantom; use serde::{Deserialize, Serialize}; @@ -60,7 +59,7 @@ impl VmExecutionExtension for HintsExtension { } } -impl VmCircuitExtension for HintsExtension { +impl VmCircuitExtension for HintsExtension { fn extend_circuit(&self, _: &mut AirInventory) -> Result<(), AirInventoryError> { Ok(()) } diff --git a/openvm-riscv/extensions/hints-transpiler/src/lib.rs b/openvm-riscv/extensions/hints-transpiler/src/lib.rs index 2c44af45c2..a28c9b0d7d 100644 --- a/openvm-riscv/extensions/hints-transpiler/src/lib.rs +++ b/openvm-riscv/extensions/hints-transpiler/src/lib.rs @@ -58,7 +58,7 @@ impl TranspilerExtension for HintsTranspilerExtension { let instruction = Instruction::phantom( PhantomDiscriminant(disc as u16), - F::from_canonical_usize(RV32_REGISTER_NUM_LIMBS * insn.rs1), + F::from_usize(RV32_REGISTER_NUM_LIMBS * insn.rs1), F::ZERO, 0, ); diff --git a/openvm-riscv/guest-ecc-manual/Cargo.toml b/openvm-riscv/guest-ecc-manual/Cargo.toml index 158c00a767..a41fe65f62 100644 --- a/openvm-riscv/guest-ecc-manual/Cargo.toml +++ b/openvm-riscv/guest-ecc-manual/Cargo.toml @@ -6,12 +6,12 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } -openvm-ecc-guest = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", subdirectory = "extensions/ecc/guest", default-features = false } -openvm-algebra-guest = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", subdirectory = "extensions/algebra/guest", default-features = false } -openvm-k256 = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", subdirectory = "guest-libs/k256", package = "k256", features = [ +openvm-ecc-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", subdirectory = "extensions/ecc/guest", default-features = false } +openvm-algebra-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", subdirectory = "extensions/algebra/guest", default-features = false } +openvm-k256 = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", subdirectory = "guest-libs/k256", package = "k256", features = [ "ecdsa", ] } diff --git a/openvm-riscv/guest-ecc-powdr-affine-hint/Cargo.toml b/openvm-riscv/guest-ecc-powdr-affine-hint/Cargo.toml index 5d0af7aba9..b5e69c78b2 100644 --- a/openvm-riscv/guest-ecc-powdr-affine-hint/Cargo.toml +++ b/openvm-riscv/guest-ecc-powdr-affine-hint/Cargo.toml @@ -5,10 +5,10 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } -k256 = { git = "https://github.com/powdr-labs/elliptic-curves-k256", rev = "a48ad5c", default-features = false, features = [ +k256 = { git = "https://github.com/powdr-labs/elliptic-curves-k256", rev = "5ee3c22", default-features = false, features = [ "expose-field", "arithmetic", ] } diff --git a/openvm-riscv/guest-ecc-projective/Cargo.toml b/openvm-riscv/guest-ecc-projective/Cargo.toml index 222fa3374d..b3aefa16ad 100644 --- a/openvm-riscv/guest-ecc-projective/Cargo.toml +++ b/openvm-riscv/guest-ecc-projective/Cargo.toml @@ -5,7 +5,7 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } k256 = { version = "0.13", default-features = false, features = ["arithmetic"] } diff --git a/openvm-riscv/guest-ecrecover-manual/Cargo.toml b/openvm-riscv/guest-ecrecover-manual/Cargo.toml index 81cbac3415..efd392c1dc 100644 --- a/openvm-riscv/guest-ecrecover-manual/Cargo.toml +++ b/openvm-riscv/guest-ecrecover-manual/Cargo.toml @@ -5,14 +5,14 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } -openvm-algebra-guest = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-algebra-moduli-macros = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-ecc-guest = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-ecc-sw-macros = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-k256 = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", package = "k256" } +openvm-algebra-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-algebra-moduli-macros = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-ecc-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-ecc-sw-macros = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-k256 = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", package = "k256" } elliptic-curve = { version = "0.13.8" } ecdsa = { version = "0.16.9" } diff --git a/openvm-riscv/guest-ecrecover-manual/openvm.toml b/openvm-riscv/guest-ecrecover-manual/openvm.toml index 4176a245be..fad151f764 100644 --- a/openvm-riscv/guest-ecrecover-manual/openvm.toml +++ b/openvm-riscv/guest-ecrecover-manual/openvm.toml @@ -1,7 +1,7 @@ [app_vm_config.rv32i] [app_vm_config.rv32m] [app_vm_config.io] -[app_vm_config.sha256] +[app_vm_config.sha2] [app_vm_config.modular] supported_moduli = [ diff --git a/openvm-riscv/guest-ecrecover/Cargo.toml b/openvm-riscv/guest-ecrecover/Cargo.toml index 2e5e3e1f62..7f078abc06 100644 --- a/openvm-riscv/guest-ecrecover/Cargo.toml +++ b/openvm-riscv/guest-ecrecover/Cargo.toml @@ -5,10 +5,10 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } -k256 = { git = "https://github.com/powdr-labs/elliptic-curves-k256", rev = "a48ad5c", default-features = false, features = [ +k256 = { git = "https://github.com/powdr-labs/elliptic-curves-k256", rev = "5ee3c22", default-features = false, features = [ "expose-field", "arithmetic", "ecdsa", diff --git a/openvm-riscv/guest-hints-test/Cargo.toml b/openvm-riscv/guest-hints-test/Cargo.toml index 406d938942..a763c47f17 100644 --- a/openvm-riscv/guest-hints-test/Cargo.toml +++ b/openvm-riscv/guest-hints-test/Cargo.toml @@ -7,7 +7,7 @@ edition = "2021" [dependencies] # The `rev` here must point to the same version used in the workspace. # Otherwise, there is conflict with the `powdr-openvm-hints-guest` dependency (which is part of the workspace). -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } powdr-openvm-hints-guest = { path = "../extensions/hints-guest/" } [profile.release-with-debug] diff --git a/openvm-riscv/guest-keccak-manual-precompile/Cargo.toml b/openvm-riscv/guest-keccak-manual-precompile/Cargo.toml index 2b34d1e897..b8d2a01ac6 100644 --- a/openvm-riscv/guest-keccak-manual-precompile/Cargo.toml +++ b/openvm-riscv/guest-keccak-manual-precompile/Cargo.toml @@ -7,6 +7,6 @@ edition = "2021" members = [] [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-keccak256 = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-keccak256 = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } diff --git a/openvm-riscv/guest-keccak/Cargo.toml b/openvm-riscv/guest-keccak/Cargo.toml index 5deaa85e1a..93ba42077b 100644 --- a/openvm-riscv/guest-keccak/Cargo.toml +++ b/openvm-riscv/guest-keccak/Cargo.toml @@ -5,7 +5,7 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } tiny-keccak = { version = "2.0.2", features = ["keccak"] } [profile.release-with-debug] diff --git a/openvm-riscv/guest-matmul/Cargo.toml b/openvm-riscv/guest-matmul/Cargo.toml index e4a5bae8b9..f7bf91375d 100644 --- a/openvm-riscv/guest-matmul/Cargo.toml +++ b/openvm-riscv/guest-matmul/Cargo.toml @@ -5,7 +5,7 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } [profile.release-with-debug] inherits = "release" diff --git a/openvm-riscv/guest-pairing-manual-precompile/Cargo.toml b/openvm-riscv/guest-pairing-manual-precompile/Cargo.toml index 5329fa73d4..df5b94a614 100644 --- a/openvm-riscv/guest-pairing-manual-precompile/Cargo.toml +++ b/openvm-riscv/guest-pairing-manual-precompile/Cargo.toml @@ -7,13 +7,13 @@ edition = "2021" members = [] [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } -openvm-algebra-guest = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", default-features = false } -openvm-ecc-guest = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", default-features = false } -openvm-pairing = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", default-features = false, features = [ +openvm-algebra-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", default-features = false } +openvm-ecc-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", default-features = false } +openvm-pairing = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", default-features = false, features = [ "bn254", ] } diff --git a/openvm-riscv/guest-pairing/Cargo.toml b/openvm-riscv/guest-pairing/Cargo.toml index d1b7571a9e..ef178bff26 100644 --- a/openvm-riscv/guest-pairing/Cargo.toml +++ b/openvm-riscv/guest-pairing/Cargo.toml @@ -7,7 +7,7 @@ edition = "2024" members = [] [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } diff --git a/openvm-riscv/guest-sha256-manual-precompile/Cargo.toml b/openvm-riscv/guest-sha256-manual-precompile/Cargo.toml index ecf4b69956..fd550274ad 100644 --- a/openvm-riscv/guest-sha256-manual-precompile/Cargo.toml +++ b/openvm-riscv/guest-sha256-manual-precompile/Cargo.toml @@ -7,6 +7,6 @@ edition = "2021" members = [] [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } -openvm-sha2 = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } +openvm-sha2 = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } diff --git a/openvm-riscv/guest-sha256-manual-precompile/openvm.toml b/openvm-riscv/guest-sha256-manual-precompile/openvm.toml index fecd2e8b39..fb0cbe8cd4 100644 --- a/openvm-riscv/guest-sha256-manual-precompile/openvm.toml +++ b/openvm-riscv/guest-sha256-manual-precompile/openvm.toml @@ -1,4 +1,4 @@ [app_vm_config.rv32i] [app_vm_config.rv32m] [app_vm_config.io] -[app_vm_config.sha256] \ No newline at end of file +[app_vm_config.sha2] \ No newline at end of file diff --git a/openvm-riscv/guest-sha256-manual-precompile/src/main.rs b/openvm-riscv/guest-sha256-manual-precompile/src/main.rs index ce99edd504..7a42609cb3 100644 --- a/openvm-riscv/guest-sha256-manual-precompile/src/main.rs +++ b/openvm-riscv/guest-sha256-manual-precompile/src/main.rs @@ -6,7 +6,7 @@ extern crate alloc; use core::hint::black_box; use openvm::io::{read, reveal_u32}; -use openvm_sha2::sha256; +use openvm_sha2::Sha256; openvm::entry!(main); @@ -14,7 +14,9 @@ pub fn main() { let n = read(); let mut output = black_box([0u8; 32]); for _ in 0..n { - output = sha256(&output); + let mut hasher = Sha256::new(); + hasher.update(&output); + output = hasher.finalize(); } reveal_u32(output[0] as u32, 0); diff --git a/openvm-riscv/guest-sha256/Cargo.toml b/openvm-riscv/guest-sha256/Cargo.toml index 5cb3ea5260..cc52280e33 100644 --- a/openvm-riscv/guest-sha256/Cargo.toml +++ b/openvm-riscv/guest-sha256/Cargo.toml @@ -5,7 +5,7 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } sha2 = { version = "0.10", default-features = false } digest = { version = "0.10", default-features = false } diff --git a/openvm-riscv/guest-u256-manual-precompile/Cargo.toml b/openvm-riscv/guest-u256-manual-precompile/Cargo.toml index cde6167b5b..fec5a26512 100644 --- a/openvm-riscv/guest-u256-manual-precompile/Cargo.toml +++ b/openvm-riscv/guest-u256-manual-precompile/Cargo.toml @@ -7,7 +7,10 @@ edition = "2021" members = [] [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } -openvm-ruint = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", package = "ruint" } +openvm-bigint-guest = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ + "export-intrinsics", +] } +openvm-ruint = { git = "https://github.com/openvm-org/uint.git", branch = "v1.14.0-openvm", package = "ruint" } diff --git a/openvm-riscv/guest-u256-manual-precompile/src/main.rs b/openvm-riscv/guest-u256-manual-precompile/src/main.rs index 2bf8de3a33..3b8f215f29 100644 --- a/openvm-riscv/guest-u256-manual-precompile/src/main.rs +++ b/openvm-riscv/guest-u256-manual-precompile/src/main.rs @@ -1,6 +1,7 @@ #![allow(clippy::needless_range_loop)] use core::array; +use openvm_bigint_guest as _; use openvm_ruint::aliases::U256; openvm::entry!(main); diff --git a/openvm-riscv/guest-u256/Cargo.toml b/openvm-riscv/guest-u256/Cargo.toml index 6aed86283a..5c4bf7a849 100644 --- a/openvm-riscv/guest-u256/Cargo.toml +++ b/openvm-riscv/guest-u256/Cargo.toml @@ -7,7 +7,7 @@ edition = "2021" members = [] [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4", features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2", features = [ "std", ] } ruint = "1.16" diff --git a/openvm-riscv/guest/Cargo.toml b/openvm-riscv/guest/Cargo.toml index c661753b3b..7738a68659 100644 --- a/openvm-riscv/guest/Cargo.toml +++ b/openvm-riscv/guest/Cargo.toml @@ -5,7 +5,7 @@ version = "0.0.0" edition = "2021" [dependencies] -openvm = { git = "https://github.com/powdr-labs/openvm.git", tag = "v1.4.2-powdr-rc.4" } +openvm = { git = "https://github.com/powdr-labs/openvm.git", branch = "v2-powdr-beta.2" } [profile.release-with-debug] inherits = "release" diff --git a/openvm-riscv/scripts/run_pairing.sh b/openvm-riscv/scripts/run_pairing.sh new file mode 100755 index 0000000000..245c083846 --- /dev/null +++ b/openvm-riscv/scripts/run_pairing.sh @@ -0,0 +1,64 @@ +#!/bin/bash + +# Script to collect some numbers from our OpenVM guest examples. +# Mostly for CI usage, but can be easily modified for manual tests. + +# NOTE: The script expects the python environment to be set up with the required +# dependencies. Should be run from the project root, will create a `results` +# directory. + +set -e + +SCRIPT_PATH=$(realpath "${BASH_SOURCE[0]}") +SCRIPTS_DIR=$(dirname "$SCRIPT_PATH") + +run_bench() { + guest="$1" + input="$2" + apcs="$3" + run_name="$4" + + echo "" + echo "==== ${run_name} ====" + echo "" + + mkdir -p "${run_name}" + + psrecord --include-children --interval 1 \ + --log "${run_name}"/psrecord.csv \ + --log-format csv \ + --plot "${run_name}"/psrecord.png \ + "cargo run --bin powdr_openvm_riscv -r --features \"metrics,cuda\" prove \"$guest\" --input \"$input\" --autoprecompiles \"$apcs\" --metrics \"${run_name}/metrics.json\" --recursion --apc-candidates-dir \"${run_name}\"" + + python3 "$SCRIPTS_DIR"/plot_trace_cells.py -o "${run_name}"/trace_cells.png "${run_name}"/metrics.json > "${run_name}"/trace_cells.txt + + # apc_candidates.json is only available when apcs > 0 + if [ "${apcs:-0}" -ne 0 ]; then + python3 "$SCRIPTS_DIR"/../../autoprecompiles/scripts/plot_effectiveness.py "${run_name}"/apc_candidates.json --output "${run_name}"/effectiveness.png + fi + + # Clean up some files that we don't want to to push. + rm -f "${run_name}"/apc_candidate_* +} + +# TODO: Some benchmarks are currently disabled to keep the nightly run below 6h. + +### Pairing +dir="results/pairing" +input="0" # No input + +mkdir -p "$dir" +pushd "$dir" + +#run_bench guest-pairing-manual-precompile "$input" 0 manual +run_bench guest-pairing "$input" 0 apc000 +run_bench guest-pairing "$input" 3 apc003 +run_bench guest-pairing "$input" 10 apc010 +run_bench guest-pairing "$input" 30 apc030 +run_bench guest-pairing "$input" 100 apc100 +run_bench guest-pairing "$input" 300 apc300 +run_bench guest-pairing "$input" 500 apc500 + +python3 $SCRIPTS_DIR/basic_metrics.py summary-table --csv **/metrics.json > basic_metrics.csv +python3 $SCRIPTS_DIR/basic_metrics.py plot **/metrics.json -o proof_time_breakdown.png +popd diff --git a/openvm-riscv/src/isa/mod.rs b/openvm-riscv/src/isa/mod.rs index 897e847013..52fd98a808 100644 --- a/openvm-riscv/src/isa/mod.rs +++ b/openvm-riscv/src/isa/mod.rs @@ -1,9 +1,10 @@ use std::collections::{BTreeSet, HashSet}; +use openvm_circuit::arch::VmField; use openvm_circuit::arch::{AirInventory, ChipInventoryError, VmBuilder}; use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP, VmOpcode}; use openvm_stark_backend::p3_field::PrimeField32; -use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine; +use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2CpuEngine; #[cfg(feature = "cuda")] use powdr_openvm::{ isa::OriginalGpuChipComplex, powdr_extension::trace_generator::SharedPeripheryChipsGpu, @@ -42,14 +43,14 @@ pub struct RiscvISA; pub struct OpenVmRegisterAddress(u8); // This seems trivial but it's tricky to put into powdr-openvm because of some From implementation issues. -impl From> for SpecializedExecutor { +impl From> for SpecializedExecutor { fn from(value: ExtendedVmConfigExecutor) -> Self { Self::OriginalExecutor(value) } } impl OpenVmISA for RiscvISA { - type Executor = ExtendedVmConfigExecutor; + type Executor = ExtendedVmConfigExecutor; type Config = ExtendedVmConfig; type CpuBuilder = ExtendedVmConfigCpuBuilder; #[cfg(feature = "cuda")] @@ -71,7 +72,7 @@ impl OpenVmISA for RiscvISA { config: &Self::Config, airs: AirInventory, ) -> Result { - >::create_chip_complex( + >::create_chip_complex( &ExtendedVmConfigCpuBuilder, config, airs, diff --git a/openvm-riscv/src/isa/symbolic_instruction_builder.rs b/openvm-riscv/src/isa/symbolic_instruction_builder.rs index eabf3b7793..75ef47d991 100644 --- a/openvm-riscv/src/isa/symbolic_instruction_builder.rs +++ b/openvm-riscv/src/isa/symbolic_instruction_builder.rs @@ -22,11 +22,11 @@ macro_rules! build_instr5 { ) -> Instruction { Instruction { opcode: VmOpcode::from_usize($code as usize), - a: T::from_canonical_u32(a), - b: T::from_canonical_u32(b), - c: T::from_canonical_u32(c), - d: T::from_canonical_u32(d), - e: T::from_canonical_u32(e), + a: T::from_u32(a), + b: T::from_u32(b), + c: T::from_u32(c), + d: T::from_u32(d), + e: T::from_u32(e), f: T::ZERO, g: T::ZERO, } @@ -53,11 +53,11 @@ macro_rules! alu_ops { ) -> Instruction { Instruction { opcode: VmOpcode::from_usize($code as usize), - a: T::from_canonical_u32(rd_ptr), - b: T::from_canonical_u32(rs1_ptr), - c: T::from_canonical_u32(rs2), + a: T::from_u32(rd_ptr), + b: T::from_u32(rs1_ptr), + c: T::from_u32(rs2), d: T::ONE, - e: T::from_canonical_u32(rs2_as), + e: T::from_u32(rs2_as), f: T::ZERO, g: T::ZERO, } @@ -86,13 +86,13 @@ macro_rules! ls_ops { ) -> Instruction { Instruction { opcode: VmOpcode::from_usize($code as usize), - a: T::from_canonical_u32(rd_rs2_ptr), - b: T::from_canonical_u32(rs1_ptr), - c: T::from_canonical_u32(imm), + a: T::from_u32(rd_rs2_ptr), + b: T::from_u32(rs1_ptr), + c: T::from_u32(imm), d: T::ONE, - e: T::from_canonical_u32(mem_as), - f: T::from_canonical_u32(needs_write), - g: T::from_canonical_u32(imm_sign), + e: T::from_u32(mem_as), + f: T::from_u32(needs_write), + g: T::from_u32(imm_sign), } } )+ @@ -115,14 +115,14 @@ macro_rules! branch_ops { imm: i32, ) -> Instruction { let imm = if imm >= 0 { - T::from_canonical_u32(imm as u32) + T::from_u32(imm as u32) } else { - -T::from_canonical_u32((-imm) as u32) + -T::from_u32((-imm) as u32) }; Instruction { opcode: VmOpcode::from_usize($code as usize), - a: T::from_canonical_u32(rs1_ptr), - b: T::from_canonical_u32(rs2_ptr), + a: T::from_u32(rs1_ptr), + b: T::from_u32(rs2_ptr), c: imm, d: T::ONE, e: T::ONE, diff --git a/openvm-riscv/src/isa/trace_generator/common.rs b/openvm-riscv/src/isa/trace_generator/common.rs index 2a5526a5e0..3c045e74d0 100644 --- a/openvm-riscv/src/isa/trace_generator/common.rs +++ b/openvm-riscv/src/isa/trace_generator/common.rs @@ -24,15 +24,9 @@ pub fn create_dummy_airs>( if let Some(keccak) = &config.keccak { VmCircuitExtension::extend_circuit(keccak, &mut inventory)?; } - if let Some(sha256) = &config.sha256 { + if let Some(sha256) = &config.sha2 { VmCircuitExtension::extend_circuit(sha256, &mut inventory)?; } - if let Some(native) = &config.native { - VmCircuitExtension::extend_circuit(native, &mut inventory)?; - } - if let Some(castf) = &config.castf { - VmCircuitExtension::extend_circuit(castf, &mut inventory)?; - } if let Some(rv32m) = &config.rv32m { VmCircuitExtension::extend_circuit(rv32m, &mut inventory)?; } diff --git a/openvm-riscv/src/isa/trace_generator/cpu.rs b/openvm-riscv/src/isa/trace_generator/cpu.rs index bc82dc14bf..0b3a55fe11 100644 --- a/openvm-riscv/src/isa/trace_generator/cpu.rs +++ b/openvm-riscv/src/isa/trace_generator/cpu.rs @@ -3,10 +3,9 @@ use openvm_bigint_circuit::Int256CpuProverExt; use openvm_circuit::arch::{AirInventory, ChipInventoryError, VmBuilder, VmProverExtension}; use openvm_circuit::system::SystemCpuBuilder; use openvm_keccak256_circuit::Keccak256CpuProverExt; -use openvm_native_circuit::NativeCpuProverExt; use openvm_pairing_circuit::PairingProverExt; use openvm_rv32im_circuit::Rv32ImCpuProverExt; -use openvm_sha256_circuit::Sha2CpuProverExt; +use openvm_sha2_circuit::Sha2CpuProverExt; use powdr_openvm::powdr_extension::trace_generator::cpu::SharedPeripheryChipsCpuProverExt; use powdr_openvm::powdr_extension::trace_generator::{DummyChipComplex, SharedPeripheryChipsCpu}; use powdr_openvm::BabyBearSC; @@ -14,7 +13,7 @@ use powdr_openvm::BabyBearSC; use crate::{ExtendedVmConfig, RiscvISA}; use openvm_ecc_circuit::EccCpuProverExt; -use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine; +use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2CpuEngine; pub fn create_dummy_chip_complex_cpu( config: &ExtendedVmConfig, @@ -22,7 +21,7 @@ pub fn create_dummy_chip_complex_cpu( shared_chips: SharedPeripheryChipsCpu, ) -> Result, ChipInventoryError> { let config = config.sdk.to_inner(); - let mut chip_complex = VmBuilder::::create_chip_complex( + let mut chip_complex = VmBuilder::::create_chip_complex( &SystemCpuBuilder, &config.system, circuit, @@ -30,7 +29,7 @@ pub fn create_dummy_chip_complex_cpu( let inventory = &mut chip_complex.inventory; // CHANGE: inject the periphery chips so that they are not created by the extensions. This is done for memory footprint: the dummy periphery chips are thrown away anyway, so we reuse a single one for all APCs. - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &SharedPeripheryChipsCpuProverExt, &shared_chips, inventory, @@ -38,84 +37,70 @@ pub fn create_dummy_chip_complex_cpu( // END CHANGE if let Some(rv32i) = &config.rv32i { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &Rv32ImCpuProverExt, rv32i, inventory, )?; } if let Some(io) = &config.io { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &Rv32ImCpuProverExt, io, inventory, )?; } if let Some(keccak) = &config.keccak { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &Keccak256CpuProverExt, keccak, inventory, )?; } - if let Some(sha256) = &config.sha256 { - VmProverExtension::::extend_prover( + if let Some(sha256) = &config.sha2 { + VmProverExtension::::extend_prover( &Sha2CpuProverExt, sha256, inventory, )?; } - if let Some(native) = &config.native { - VmProverExtension::::extend_prover( - &NativeCpuProverExt, - native, - inventory, - )?; - } - if let Some(castf) = &config.castf { - VmProverExtension::::extend_prover( - &NativeCpuProverExt, - castf, - inventory, - )?; - } if let Some(rv32m) = &config.rv32m { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &Rv32ImCpuProverExt, rv32m, inventory, )?; } if let Some(bigint) = &config.bigint { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &Int256CpuProverExt, bigint, inventory, )?; } if let Some(modular) = &config.modular { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &AlgebraCpuProverExt, modular, inventory, )?; } if let Some(fp2) = &config.fp2 { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &AlgebraCpuProverExt, fp2, inventory, )?; } if let Some(pairing) = &config.pairing { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &PairingProverExt, pairing, inventory, )?; } if let Some(ecc) = &config.ecc { - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &EccCpuProverExt, ecc, inventory, diff --git a/openvm-riscv/src/isa/trace_generator/cuda.rs b/openvm-riscv/src/isa/trace_generator/cuda.rs index 8aa20905a7..d91aba968d 100644 --- a/openvm-riscv/src/isa/trace_generator/cuda.rs +++ b/openvm-riscv/src/isa/trace_generator/cuda.rs @@ -7,7 +7,7 @@ use powdr_openvm::{ powdr_extension::trace_generator::cuda::{ GpuDummyChipComplex, SharedPeripheryChipsGpu, SharedPeripheryChipsGpuProverExt, }, - BabyBearSC, GpuBabyBearPoseidon2Engine, + BabyBearSC, GpuBabyBearPoseidon2CpuEngine, }; use crate::{ExtendedVmConfig, RiscvISA}; @@ -21,11 +21,10 @@ pub fn create_dummy_chip_complex_gpu( use openvm_bigint_circuit::Int256GpuProverExt; use openvm_ecc_circuit::EccProverExt; use openvm_keccak256_circuit::Keccak256GpuProverExt; - use openvm_native_circuit::NativeGpuProverExt; use openvm_rv32im_circuit::Rv32ImGpuProverExt; - use openvm_sha256_circuit::Sha256GpuProverExt; + use openvm_sha2_circuit::Sha2GpuProverExt; - type E = GpuBabyBearPoseidon2Engine; + type E = GpuBabyBearPoseidon2CpuEngine; let config = config.sdk.to_inner(); let mut chip_complex = @@ -49,14 +48,8 @@ pub fn create_dummy_chip_complex_gpu( if let Some(keccak) = &config.keccak { VmProverExtension::::extend_prover(&Keccak256GpuProverExt, keccak, inventory)?; } - if let Some(sha256) = &config.sha256 { - VmProverExtension::::extend_prover(&Sha256GpuProverExt, sha256, inventory)?; - } - if let Some(native) = &config.native { - VmProverExtension::::extend_prover(&NativeGpuProverExt, native, inventory)?; - } - if let Some(castf) = &config.castf { - VmProverExtension::::extend_prover(&NativeGpuProverExt, castf, inventory)?; + if let Some(sha2) = &config.sha2 { + VmProverExtension::::extend_prover(&Sha2GpuProverExt, sha2, inventory)?; } if let Some(rv32m) = &config.rv32m { VmProverExtension::::extend_prover(&Rv32ImGpuProverExt, rv32m, inventory)?; diff --git a/openvm-riscv/src/lib.rs b/openvm-riscv/src/lib.rs index 667f329960..d84ed9d6cc 100644 --- a/openvm-riscv/src/lib.rs +++ b/openvm-riscv/src/lib.rs @@ -5,28 +5,27 @@ use eyre::Result; use openvm_build::{build_guest_package, find_unique_executable, get_package, TargetFilter}; +#[cfg(debug_assertions)] +use openvm_circuit::arch::debug_proving_ctx; use openvm_circuit::arch::execution_mode::metered::segment_ctx::SegmentationLimits; #[cfg(feature = "cuda")] use openvm_circuit::arch::DenseRecordArena; use openvm_circuit::arch::{ - debug_proving_ctx, AirInventory, ChipInventoryError, InitFileGenerator, MatrixRecordArena, - SystemConfig, VmBuilder, VmChipComplex, VmProverExtension, + AirInventory, ChipInventoryError, InitFileGenerator, MatrixRecordArena, SystemConfig, + VmBuilder, VmChipComplex, VmProverExtension, }; #[cfg(feature = "cuda")] use openvm_circuit::system::cuda::SystemChipInventoryGPU; use openvm_circuit::system::SystemChipInventory; -use openvm_sdk::config::SdkVmCpuBuilder; +use openvm_sdk_config::{SdkVmConfig, SdkVmConfigExecutor, SdkVmCpuBuilder, TranspilerConfig}; -use openvm_sdk::config::TranspilerConfig; -use openvm_sdk::prover::{verify_app_proof, AggStarkProver}; +use openvm_cpu_backend::{CpuBackend, CpuDevice}; use openvm_sdk::{ - config::{AppConfig, SdkVmConfig, SdkVmConfigExecutor, DEFAULT_APP_LOG_BLOWUP}, - Sdk, StdIn, + config::{AggregationSystemParams, AppConfig}, + StdIn, }; -use openvm_stark_backend::config::Val; -use openvm_stark_backend::engine::StarkEngine; -use openvm_stark_backend::prover::cpu::{CpuBackend, CpuDevice}; -use openvm_stark_sdk::config::FriParameters; +use openvm_stark_backend::{StarkEngine, Val}; +use openvm_stark_sdk::config::{app_params_with_100_bits_security, MAX_APP_LOG_STACKED_HEIGHT}; use openvm_stark_sdk::p3_baby_bear::BabyBear; use openvm_transpiler::transpiler::Transpiler; use powdr_autoprecompiles::empirical_constraints::EmpiricalConstraints; @@ -39,7 +38,7 @@ use powdr_openvm::BabyBearSC; #[cfg(not(feature = "cuda"))] use powdr_openvm::PowdrSdkCpu; #[cfg(feature = "cuda")] -use powdr_openvm::{GpuBabyBearPoseidon2Engine, GpuBackend, PowdrSdkGpu}; +use powdr_openvm::{GpuBabyBearPoseidon2CpuEngine, GpuBackend, PowdrSdkGpu}; use powdr_openvm_riscv_hints_circuit::{HintsExtension, HintsExtensionExecutor, HintsProverExt}; use powdr_openvm_riscv_hints_transpiler::HintsTranspilerExtension; use serde::{Deserialize, Serialize}; @@ -104,23 +103,31 @@ pub fn compile_openvm( // try to load the sdk config from the openvm.toml file, otherwise use the default let openvm_toml_path = path.join("openvm.toml"); - let app_config = if openvm_toml_path.exists() { - let toml = std::fs::read_to_string(&openvm_toml_path)?; - toml::from_str(&toml)? + let system_params = app_params_with_100_bits_security(MAX_APP_LOG_STACKED_HEIGHT); + let app_config: AppConfig = if openvm_toml_path.exists() { + let toml_str = std::fs::read_to_string(&openvm_toml_path)?; + // Deserialize just the app_vm_config from the TOML, then pair with our system_params. + // The TOML files don't contain system_params (v2 addition). + #[derive(serde::Deserialize)] + struct PartialAppConfig { + app_vm_config: SdkVmConfig, + } + let partial: PartialAppConfig = toml::from_str(&toml_str)?; + AppConfig::new(partial.app_vm_config, system_params) } else { - AppConfig::riscv32() + AppConfig::riscv32(system_params) }; - let mut sdk = Sdk::new(app_config)?; - - let transpiler = sdk.transpiler().unwrap(); - - // Add our custom transpiler extensions - sdk.set_transpiler( - transpiler - .clone() - .with_extension(HintsTranspilerExtension {}), - ); + // Build SDK with our custom transpiler extension + let transpiler = app_config + .app_vm_config + .transpiler() + .with_extension(HintsTranspilerExtension {}); + let sdk = openvm_sdk::Sdk::builder() + .app_config(app_config) + .agg_params(AggregationSystemParams::default()) + .transpiler(transpiler) + .build_without_transpiler()?; let elf = sdk.build( guest_opts.clone(), @@ -242,7 +249,7 @@ where pub struct ExtendedVmConfigGpuBuilder; #[cfg(feature = "cuda")] -impl VmBuilder for ExtendedVmConfigGpuBuilder { +impl VmBuilder for ExtendedVmConfigGpuBuilder { type VmConfig = ExtendedVmConfig; type SystemChipInventory = SystemChipInventoryGPU; type RecordArena = DenseRecordArena; @@ -255,13 +262,13 @@ impl VmBuilder for ExtendedVmConfigGpuBuilder { VmChipComplex, ChipInventoryError, > { - let mut chip_complex = VmBuilder::::create_chip_complex( - &openvm_sdk::config::SdkVmGpuBuilder, + let mut chip_complex = VmBuilder::::create_chip_complex( + &openvm_sdk_config::SdkVmGpuBuilder, &config.sdk, circuit, )?; let inventory = &mut chip_complex.inventory; - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &HintsProverExt, &config.hints, inventory, @@ -291,61 +298,69 @@ pub fn prove( inputs: StdIn, segment_height: Option, // uses the default height if None ) -> Result<(), Box> { + // Apply the segment-height override up-front so the mock path (which re-reads + // `program.vm_config` via `do_with_trace`) also benefits. Without this, only the + // proving path below was shrinking per-chip GPU allocations — mock runs still + // blew past device memory on shared GPU runners. + let mut program_storage; + let program = if let Some(segment_height) = segment_height { + program_storage = program.clone(); + program_storage + .vm_config + .original + .config_mut() + .sdk + .system + .config + .segmentation_config + .limits = SegmentationLimits::default().with_max_trace_height(segment_height as u32); + tracing::debug!("Setting max segment len to {}", segment_height); + &program_storage + } else { + program + }; + if mock { - do_with_trace(program, inputs, |_segment_idx, vm, pk, ctx| { - debug_proving_ctx(vm, pk, &ctx); + do_with_trace(program, inputs, |_segment_idx, _vm, _pk, _ctx| { + #[cfg(debug_assertions)] + debug_proving_ctx(_vm, &_ctx); + #[cfg(not(debug_assertions))] + tracing::warn!("mock proving skips debug checks in release mode"); })?; } else { let exe = &program.exe; - let mut vm_config = program.vm_config.clone(); - - // DefaultSegmentationStrategy { max_segment_len: 4194204, max_cells_per_chip_in_segment: 503304480 } - if let Some(segment_height) = segment_height { - vm_config - .original - .config_mut() - .sdk - .system - .config - .segmentation_limits = - SegmentationLimits::default().with_max_trace_height(segment_height as u32); - tracing::debug!("Setting max segment len to {}", segment_height); - } + let vm_config = program.vm_config.clone(); // Set app configuration - let app_fri_params = - FriParameters::standard_with_100_bits_conjectured_security(DEFAULT_APP_LOG_BLOWUP); - let app_config = AppConfig::new(app_fri_params, vm_config.clone()); + let system_params = app_params_with_100_bits_security(MAX_APP_LOG_STACKED_HEIGHT); + let app_config = AppConfig::new(vm_config.clone(), system_params); // Create the SDK #[cfg(feature = "cuda")] - let sdk = PowdrSdkGpu::new(app_config).unwrap(); + let sdk = PowdrSdkGpu::new(app_config, AggregationSystemParams::default()).unwrap(); #[cfg(not(feature = "cuda"))] - let sdk = PowdrSdkCpu::new(app_config).unwrap(); - let mut app_prover = sdk.app_prover(exe.clone())?; - - // Generate a proof - tracing::info!("Generating app proof..."); - let start = std::time::Instant::now(); - let app_proof = app_prover.prove(inputs.clone())?; - tracing::info!("App proof took {:?}", start.elapsed()); - - tracing::info!("Public values: {:?}", app_proof.user_public_values); - - // Verify - let app_vk = sdk.app_pk().get_app_vk(); - verify_app_proof(&app_vk, &app_proof)?; - tracing::info!("App proof verification done."); + let sdk = PowdrSdkCpu::new(app_config, AggregationSystemParams::default()).unwrap(); if recursion { - let mut agg_prover: AggStarkProver<_, _> = sdk.prover(exe.clone())?.agg_prover; + // Use StarkProver which does app proving + aggregation (recursion) + let mut stark_prover = sdk.prover(exe.clone())?; + tracing::info!("Generating STARK proof (app + aggregation)..."); + let start = std::time::Instant::now(); + let _stark_proof = stark_prover.prove(inputs.clone(), &[])?; + tracing::info!("STARK proof (with recursion) took {:?}", start.elapsed()); + } else { + let mut app_prover = sdk.app_prover(exe.clone())?; - // Note that this proof is not verified. We assume that any valid app proof - // (verified above) also leads to a valid aggregation proof. - // If this was not the case, it would be a completeness bug in OpenVM. + // Generate a proof + tracing::info!("Generating app proof..."); let start = std::time::Instant::now(); - let _ = agg_prover.generate_root_verifier_input(app_proof)?; - tracing::info!("Agg proof (inner recursion) took {:?}", start.elapsed()); + let app_proof = app_prover.prove(inputs.clone())?; + tracing::info!("App proof took {:?}", start.elapsed()); + + tracing::info!("Public values: {:?}", app_proof.user_public_values); + + // Note: verification is done automatically in debug_assertions mode inside prove() + tracing::info!("App proof generation done."); } tracing::info!("All done."); @@ -367,6 +382,13 @@ mod tests { use pretty_assertions::assert_eq; use test_log::test; + // Cap per-chip GPU trace allocations for tests that don't care about the + // default segment height. The OpenVM default is 2^22 rows (~4 GiB per chip), + // which pushes concurrent GPU tests past the memory on our shared runner. + // All small test inputs here fit easily in 2^20 rows and still run in a + // single segment. + const TEST_DEFAULT_SEGMENT_HEIGHT: usize = 1 << 20; + #[allow(clippy::too_many_arguments)] fn compile_and_prove( guest: &str, @@ -380,6 +402,7 @@ mod tests { let guest = compile_openvm(guest, GuestOptions::default()).unwrap(); let program = compile_exe(guest, config, pgo_config, EmpiricalConstraints::default()).unwrap(); + let segment_height = segment_height.or(Some(TEST_DEFAULT_SEGMENT_HEIGHT)); prove(&program, mock, recursion, stdin, segment_height) } @@ -583,8 +606,8 @@ mod tests { let mut stdin = StdIn::default(); stdin.write(&GUEST_KECCAK_ITER_SMALL); let config = default_powdr_openvm_config(GUEST_KECCAK_APC, GUEST_KECCAK_SKIP); - // should create two segments - prove_simple(GUEST_KECCAK, config, stdin, PgoConfig::None, Some(4_000)); + // should create two segments (v2 requires power-of-two max_trace_height) + prove_simple(GUEST_KECCAK, config, stdin, PgoConfig::None, Some(4_096)); } #[test] @@ -1111,11 +1134,10 @@ mod tests { const NON_POWDR_EXPECTED_MACHINE_COUNT: usize = 19; const NON_POWDR_EXPECTED_SUM: AirMetrics = AirMetrics { widths: AirWidths { - preprocessed: 7, - main: 798, - log_up: 684, + preprocessed: 0, + main: 819, }, - constraints: 604, + constraints: 643, bus_interactions: 253, }; @@ -1139,7 +1161,6 @@ mod tests { widths: AirWidths { preprocessed: 0, main: 38, - log_up: 56, }, constraints: 12, bus_interactions: 26, @@ -1167,7 +1188,6 @@ mod tests { widths: AirWidths { preprocessed: 0, main: 38, - log_up: 56, }, constraints: 12, bus_interactions: 26, @@ -1184,12 +1204,10 @@ mod tests { before: AirWidths { preprocessed: 0, main: 170, - log_up: 236, }, after: AirWidths { preprocessed: 0, main: 38, - log_up: 56, }, } "#]]), @@ -1216,7 +1234,6 @@ mod tests { widths: AirWidths { preprocessed: 0, main: 14254, - log_up: 22752, }, constraints: 4279, bus_interactions: 11143, @@ -1244,7 +1261,6 @@ mod tests { widths: AirWidths { preprocessed: 0, main: 14226, - log_up: 22720, }, constraints: 4255, bus_interactions: 11133, @@ -1261,12 +1277,10 @@ mod tests { before: AirWidths { preprocessed: 0, main: 183410, - log_up: 227144, }, after: AirWidths { preprocessed: 0, main: 14226, - log_up: 22720, }, } "#]]), @@ -1292,11 +1306,10 @@ mod tests { AirMetrics { widths: AirWidths { preprocessed: 0, - main: 17184, - log_up: 27796, + main: 16657, }, - constraints: 8573, - bus_interactions: 11892, + constraints: 8357, + bus_interactions: 11554, } "#]], powdr_expected_machine_count: expect![[r#" @@ -1309,13 +1322,11 @@ mod tests { AirWidthsDiff { before: AirWidths { preprocessed: 0, - main: 127688, - log_up: 169860, + main: 125426, }, after: AirWidths { preprocessed: 0, - main: 17184, - log_up: 27796, + main: 16657, }, } "#]]), @@ -1341,11 +1352,10 @@ mod tests { AirMetrics { widths: AirWidths { preprocessed: 0, - main: 19873, - log_up: 30884, + main: 18508, }, - constraints: 10968, - bus_interactions: 13423, + constraints: 10511, + bus_interactions: 12483, } "#]], powdr_expected_machine_count: expect![[r#" @@ -1358,13 +1368,11 @@ mod tests { AirWidthsDiff { before: AirWidths { preprocessed: 0, - main: 150546, - log_up: 198172, + main: 142768, }, after: AirWidths { preprocessed: 0, - main: 19873, - log_up: 30884, + main: 18508, }, } "#]]), @@ -1391,7 +1399,6 @@ mod tests { widths: AirWidths { preprocessed: 0, main: 2022, - log_up: 3472, }, constraints: 187, bus_interactions: 1734, @@ -1419,7 +1426,6 @@ mod tests { widths: AirWidths { preprocessed: 0, main: 2022, - log_up: 3472, }, constraints: 187, bus_interactions: 1734, @@ -1447,7 +1453,6 @@ mod tests { widths: AirWidths { preprocessed: 0, main: 2022, - log_up: 3472, }, constraints: 187, bus_interactions: 1734, @@ -1464,12 +1469,10 @@ mod tests { before: AirWidths { preprocessed: 0, main: 27521, - log_up: 35156, }, after: AirWidths { preprocessed: 0, main: 2022, - log_up: 3472, }, } "#]]), @@ -1498,15 +1501,14 @@ mod tests { AirMetrics { widths: AirWidths { preprocessed: 0, - main: 3234, - log_up: 5264, + main: 6992, }, - constraints: 571, - bus_interactions: 2562, + constraints: 1706, + bus_interactions: 5180, } "#]], powdr_expected_machine_count: expect![[r#" - 22 + 61 "#]], non_powdr_expected_sum: NON_POWDR_EXPECTED_SUM, non_powdr_expected_machine_count: NON_POWDR_EXPECTED_MACHINE_COUNT, @@ -1515,13 +1517,11 @@ mod tests { AirWidthsDiff { before: AirWidths { preprocessed: 0, - main: 32376, - log_up: 41660, + main: 46806, }, after: AirWidths { preprocessed: 0, - main: 3234, - log_up: 5264, + main: 6992, }, } "#]]), @@ -1538,7 +1538,7 @@ mod tests { } mod extraction { - use crate::{ExtendedVmConfig, RiscvISA, DEFAULT_OPENVM_DEGREE_BOUND}; + use crate::{ExtendedVmConfig, RiscvISA}; use openvm_algebra_circuit::{Fp2Extension, ModularExtension}; use openvm_bigint_circuit::Int256; @@ -1546,7 +1546,7 @@ mod tests { use openvm_ecc_circuit::{WeierstrassExtension, SECP256K1_CONFIG}; use openvm_pairing_circuit::{PairingCurve, PairingExtension}; use openvm_rv32im_circuit::Rv32M; - use openvm_sdk::config::SdkVmConfig; + use openvm_sdk_config::SdkVmConfig; use powdr_openvm::extraction_utils::OriginalVmConfig; use powdr_openvm_riscv_hints_circuit::HintsExtension; @@ -1554,10 +1554,7 @@ mod tests { fn test_get_bus_map() { let use_kzg_intrinsics = true; - let system_config = SystemConfig::default() - .with_continuations() - .with_max_constraint_degree(DEFAULT_OPENVM_DEGREE_BOUND) - .with_public_values(32); + let system_config = SystemConfig::default().with_public_values(32); let int256 = Int256::default(); let bn_config = PairingCurve::Bn254.curve_config(); let bls_config = PairingCurve::Bls12_381.curve_config(); @@ -1588,7 +1585,7 @@ mod tests { .rv32m(rv32m) .io(Default::default()) .keccak(Default::default()) - .sha256(Default::default()) + .sha2(Default::default()) .bigint(int256) .modular(ModularExtension::new(supported_moduli)) .fp2(Fp2Extension::new(supported_complex_moduli)) diff --git a/openvm-riscv/tests/common/mod.rs b/openvm-riscv/tests/common/mod.rs index 7ab76dc394..1a11d99586 100644 --- a/openvm-riscv/tests/common/mod.rs +++ b/openvm-riscv/tests/common/mod.rs @@ -1,5 +1,5 @@ use openvm_instructions::instruction::Instruction; -use openvm_sdk::config::SdkVmConfig; +use openvm_sdk_config::SdkVmConfig; use openvm_stark_sdk::p3_baby_bear::BabyBear; use powdr_autoprecompiles::blocks::SuperBlock; use powdr_openvm::extraction_utils::OriginalVmConfig; diff --git a/openvm/Cargo.toml b/openvm/Cargo.toml index dc225fec52..da34d7df72 100644 --- a/openvm/Cargo.toml +++ b/openvm/Cargo.toml @@ -19,6 +19,7 @@ itertools.workspace = true openvm-circuit.workspace = true openvm-circuit-primitives.workspace = true openvm-instructions.workspace = true +openvm-cpu-backend.workspace = true openvm-stark-backend.workspace = true openvm-stark-sdk.workspace = true powdr-autoprecompiles.workspace = true @@ -27,10 +28,10 @@ powdr-openvm-bus-interaction-handler.workspace = true openvm-circuit-derive.workspace = true openvm-circuit-primitives-derive.workspace = true openvm-sdk.workspace = true +openvm-sdk-config.workspace = true openvm-cuda-backend = { workspace = true, optional = true } openvm-cuda-common = { workspace = true, optional = true } openvm-transpiler.workspace = true -openvm-native-circuit.workspace = true serde.workspace = true powdr-number.workspace = true derive_more.workspace = true @@ -42,6 +43,7 @@ indicatif = "0.18.3" cfg-if = "1.0.4" powdr-riscv-elf.workspace = true pretty_assertions = { workspace = true, optional = true } +rayon = "1.10" [build-dependencies] openvm-cuda-builder = { workspace = true, optional = true } diff --git a/openvm/build.rs b/openvm/build.rs index 808e715b13..3658c9f85a 100644 --- a/openvm/build.rs +++ b/openvm/build.rs @@ -14,6 +14,7 @@ fn main() { .include("cuda/include") // Point to header file folder of this crate .watch("cuda") // Watch file changes of this crate for recompilation .library_name("powdr_gpu") // Library name of this crate; doesn't affect import name + .link_lib("nvrtc") // NVRTC for runtime CUDA compilation .files_from_glob("cuda/src/**/*.cu"); // Import all `.cu` files with zero or more nested sub-folders under `cuda/src`of this crate builder.emit_link_directives(); diff --git a/openvm/cuda/src/apc_apply_bus.cu b/openvm/cuda/src/apc_apply_bus.cu index 646c3b276e..1d4417dbee 100644 --- a/openvm/cuda/src/apc_apply_bus.cu +++ b/openvm/cuda/src/apc_apply_bus.cu @@ -75,7 +75,7 @@ __global__ void apc_apply_bus_kernel( uint32_t value = v_fp.asUInt32(); uint32_t max_bits = b_fp.asUInt32(); lookup::Histogram hist(d_var_hist, (uint32_t)var_num_bins); - uint32_t idx = (1u << max_bits) + value; // `max_bit` + uint32_t idx = (1u << max_bits) + value - 1u; // matches VariableRangeChecker::add_count // apply multiplicity by looping; warp-level dedup in Histogram minimizes contention for (uint32_t k = 0; k < (uint32_t)mult.asUInt32(); ++k) hist.add_count(idx); @@ -121,6 +121,101 @@ __global__ void apc_apply_bus_kernel( } } +// ============================================================================================ +// Phase 0 spike kernel — same structure, no bytecode evaluation. +// Used to determine whether the bus_kernel's ~76 ms cost is bytecode-dispatch +// dominated (in which case NVRTC will help) or atomic-traffic dominated (in +// which case NVRTC won't help and we should pivot). +// +// Each (interaction, row) pair issues ONE trace-cell load (to keep DRAM +// traffic representative) and ONE atomicAdd to the appropriate periphery +// histogram. Index/value derived cheaply from the trace cell so the +// compiler can't optimize out the load. mult is treated as 1 (typical for +// Receive interactions). +// ============================================================================================ +__global__ void apc_apply_bus_kernel_spike( + const Fp* __restrict__ d_output, + int num_apc_calls, + size_t n_interactions, + const DevInteraction* __restrict__ d_interactions, + uint32_t var_range_bus_id, + uint32_t* __restrict__ d_var_hist, + size_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 int warp = (threadIdx.x >> 5); + const int lane = (threadIdx.x & 31); + const int warps_per_block = (blockDim.x >> 5); + + const int H = num_apc_calls; // assume H == num_apc_calls for stride; close enough for measurement + + for (int i = blockIdx.x * warps_per_block + warp; i < (int)n_interactions; i += gridDim.x * warps_per_block) { + DevInteraction intr = d_interactions[i]; + + for (int r = lane; r < num_apc_calls; r += 32) { + // Load ONE trace cell to mimic DRAM traffic. Cell choice: column 0 of + // current row (any deterministic choice; we just want a memory load). + Fp cell = d_output[(size_t)0 * (size_t)H + (size_t)r]; + uint32_t cell_u = cell.asUInt32(); + + if (intr.bus_id == var_range_bus_id) { + lookup::Histogram hist(d_var_hist, (uint32_t)var_num_bins); + // Index derived from cell + interaction index, masked to bin count. + uint32_t idx = (cell_u + (uint32_t)i) % (uint32_t)var_num_bins; + hist.add_count(idx); + } else if (intr.bus_id == tuple2_bus_id) { + uint32_t total = tuple2_sz0 * tuple2_sz1; + if (total > 0) { + lookup::Histogram hist(d_tuple2_hist, total); + uint32_t idx = (cell_u + (uint32_t)i) % total; + hist.add_count(idx); + } + } else if (intr.bus_id == bitwise_bus_id) { + BitwiseOperationLookup bl(d_bitwise_hist, BITWISE_NUM_BITS); + // Use low byte twice; selector range path (no xor) for simplicity. + bl.add_range(cell_u & 0xFFu, (cell_u >> 8) & 0xFFu); + } + } + } +} + +extern "C" int _apc_apply_bus_spike( + const Fp* d_output, + int num_apc_calls, + const DevInteraction* d_interactions, + size_t n_interactions, + uint32_t var_range_bus_id, + uint32_t* d_var_hist, + size_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 +) { + const int block_x = 256; + const dim3 block(block_x, 1, 1); + const unsigned warps_per_block = (unsigned)(block_x / 32); + size_t g_size = (n_interactions + (size_t)warps_per_block - 1) / (size_t)warps_per_block; + unsigned g = (unsigned)g_size; + if (g == 0u) g = 1u; + const dim3 grid(g, 1, 1); + + apc_apply_bus_kernel_spike<<>>( + d_output, num_apc_calls, n_interactions, d_interactions, + 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 + ); + return (int)cudaGetLastError(); +} + // ============================================================================================ // Host launcher wrapper — callable from Rust FFI or cudarc // ============================================================================================ diff --git a/openvm/cuda/src/apc_jit_tracegen.cu b/openvm/cuda/src/apc_jit_tracegen.cu new file mode 100644 index 0000000000..145b05d579 --- /dev/null +++ b/openvm/cuda/src/apc_jit_tracegen.cu @@ -0,0 +1,566 @@ +#include "primitives/buffer_view.cuh" +#include "primitives/constants.h" +#include "primitives/trace_access.h" +#include "expr_eval.cuh" + +// ============================================================================================ +// JIT column computation descriptors +// ============================================================================================ + +// Computation type IDs — must match Rust-side ColumnComputation enum +enum JitCompType : uint16_t { + JIT_DIRECT_U32 = 0, + JIT_DIRECT_U8 = 1, + JIT_DIRECT_U16 = 2, + JIT_TIMESTAMP_DECOMP = 3, + JIT_ALU_RESULT = 4, + JIT_BOOL_FROM_OPCODE = 5, + JIT_POINTER_LIMB = 6, + JIT_SHIFT_RESULT = 7, + JIT_SHIFT_BIT_MUL_L = 8, + JIT_SHIFT_BIT_MUL_R = 9, + JIT_SHIFT_B_SIGN = 10, + JIT_SHIFT_BIT_MARKER = 11, + JIT_SHIFT_LIMB_MARKER = 12, + JIT_SHIFT_BIT_CARRY = 13, + JIT_BRANCH_EQ_CMP = 14, + JIT_BRANCH_EQ_DIFF_INV= 15, + JIT_CONSTANT = 16, + JIT_LS_RD_RS2_PTR = 17, + JIT_LS_NEEDS_WRITE = 18, + JIT_LS_WRITE_AUX_TS = 19, + JIT_LS_WRITE_AUX_DEC = 20, + JIT_LS_IS_LOAD = 21, + JIT_LS_FLAG = 22, + JIT_LS_WRITE_DATA = 23, + // LessThan (slt/sltu) — share (opcode, b, c) signature for future fusion. + JIT_LT_CMP = 24, + JIT_LT_DIFF_VAL = 25, + JIT_LT_DIFF_MARKER = 26, + JIT_LT_B_MSB_F = 27, + JIT_LT_C_MSB_F = 28, + // BranchLessThan (blt/bltu/bge/bgeu) — share (opcode, a, b) signature. + JIT_BLT_CMP = 29, + JIT_BLT_CMP_LT = 30, + JIT_BLT_DIFF_VAL = 31, + JIT_BLT_DIFF_MARKER = 32, + JIT_BLT_A_MSB_F = 33, + JIT_BLT_B_MSB_F = 34, + // Auipc / Jalr arms. + JIT_AUIPC_RD_LIMB = 35, + JIT_JALR_TO_PC_LSB = 36, + JIT_JALR_TO_PC_LIMB = 37, + JIT_JALR_RD_LIMB = 38, + // Conditional flag bits — set on `comp_type` to gate the column write. + // p[5] holds the byte offset of the gate value within the record. + // JIT_COND_FLAG (0x80): record[p[5]] != 0 (1-byte gate) + // JIT_COND_NOT_MAX_U32 (0x40): u32 at p[5] != 0xFFFFFFFF (4-byte gate) + JIT_COND_FLAG = 0x80, + JIT_COND_NOT_MAX_U32 = 0x40, + JIT_COND_DIRECT_U32 = JIT_COND_FLAG | JIT_DIRECT_U32, + JIT_COND_TS_DECOMP = JIT_COND_FLAG | JIT_TIMESTAMP_DECOMP, +}; + +// Describes how to compute one APC column from record bytes. +// 16 bytes total, well-aligned. +extern "C" { +struct JitColumnDesc { + uint16_t comp_type; // JitCompType + uint16_t apc_col; // destination column in APC trace + uint16_t p[6]; // parameters (byte offsets, indices, etc.) +}; + +// Describes one instruction within the APC. +struct JitInstructionDesc { + uint32_t arena_offset; // byte offset into concatenated arena buffer + uint32_t record_stride; // bytes between records for consecutive APC calls + uint32_t record_offset; // this instruction's byte offset within one APC call's records + uint32_t col_desc_start; // index into column descriptor array + uint32_t col_desc_count; // number of column descriptors for this instruction +}; +} + +// ============================================================================================ +// Inline helpers for shift operations +// ============================================================================================ + +__device__ __forceinline__ void get_shift_amounts(uint8_t c0, int &limb_shift, int &bit_shift) { + int shift = c0 % 32; // NUM_LIMBS * CELL_BITS = 4 * 8 = 32 + limb_shift = shift / 8; + bit_shift = shift % 8; +} + +__device__ __forceinline__ void run_alu(uint8_t opcode, const uint8_t *b, const uint8_t *c, uint8_t *a) { + switch (opcode) { + case 0: { // ADD + uint32_t carry = 0; + for (int i = 0; i < 4; i++) { + uint32_t s = (uint32_t)b[i] + (uint32_t)c[i] + carry; + a[i] = s & 0xFF; carry = s >> 8; + } + break; + } + case 1: { // SUB + uint32_t borrow = 0; + for (int i = 0; i < 4; i++) { + uint32_t rhs = (uint32_t)c[i] + borrow; + if ((uint32_t)b[i] >= rhs) { a[i] = b[i] - rhs; borrow = 0; } + else { a[i] = (256 + b[i] - rhs); borrow = 1; } + } + break; + } + case 2: for (int i=0;i<4;i++) a[i] = b[i] ^ c[i]; break; // XOR + case 3: for (int i=0;i<4;i++) a[i] = b[i] | c[i]; break; // OR + case 4: for (int i=0;i<4;i++) a[i] = b[i] & c[i]; break; // AND + default: for (int i=0;i<4;i++) a[i] = 0; break; + } +} + +__device__ __forceinline__ void run_shift(uint8_t opcode, const uint8_t *b, const uint8_t *c, uint8_t *a) { + int ls, bs; + get_shift_amounts(c[0], ls, bs); + if (opcode == 0) { // SLL + for (int i=0;i<4;i++) a[i] = 0; + for (int i=ls;i<4;i++) { + uint16_t high = (uint16_t)b[i-ls] << bs; + uint16_t low = (i>ls) ? ((uint16_t)b[i-ls-1] >> (8-bs)) : 0; + a[i] = (high | low) & 0xFF; + } + } else { // SRL/SRA + uint8_t msb = b[3] >> 7; + uint8_t fill = (opcode==1) ? 0 : (0xFF * msb); + for (int i=0;i<4;i++) a[i] = fill; + for (int i=0;i<4-ls;i++) { + uint16_t p1 = (uint16_t)(b[i+ls] >> bs); + uint16_t p2v = (i+ls+1<4) ? b[i+ls+1] : fill; + uint16_t p2 = (uint16_t)p2v << (8-bs); + a[i] = (p1 | p2) & 0xFF; + } + } +} + +// Mirror of LessThanCoreAir::run_less_than (NUM_LIMBS=4, LIMB_BITS=8). +// Returns cmp_result via the function value, sets out_diff_idx (0..4 with 4 +// meaning equal), out_b_sign, out_c_sign as side outputs. +__device__ __forceinline__ bool run_less_than( + bool is_slt, const uint8_t *b, const uint8_t *c, + int *out_diff_idx, bool *out_b_sign, bool *out_c_sign +) { + bool b_sign = is_slt && ((b[3] >> 7) == 1); + bool c_sign = is_slt && ((c[3] >> 7) == 1); + *out_b_sign = b_sign; + *out_c_sign = c_sign; + for (int i = 3; i >= 0; --i) { + if (b[i] != c[i]) { + *out_diff_idx = i; + return ((b[i] < c[i]) ^ b_sign ^ c_sign); + } + } + *out_diff_idx = 4; + return false; +} + +// Mirror of BranchLessThanCoreAir::run_cmp (NUM_LIMBS=4, LIMB_BITS=8). +// Local opcodes: BLT=0, BLTU=1, BGE=2, BGEU=3. +__device__ __forceinline__ bool run_branch_lt( + uint8_t local_opcode, const uint8_t *a, const uint8_t *b, + int *out_diff_idx, bool *out_a_sign, bool *out_b_sign +) { + bool signed_op = (local_opcode == 0) || (local_opcode == 2); + bool ge_op = (local_opcode == 2) || (local_opcode == 3); + bool a_sign = signed_op && ((a[3] >> 7) == 1); + bool b_sign = signed_op && ((b[3] >> 7) == 1); + *out_a_sign = a_sign; + *out_b_sign = b_sign; + for (int i = 3; i >= 0; --i) { + if (a[i] != b[i]) { + *out_diff_idx = i; + return ((a[i] < b[i]) ^ a_sign ^ b_sign ^ ge_op); + } + } + *out_diff_idx = 4; + return ge_op; +} + +// ============================================================================================ +// eval_jit_column: evaluates one column descriptor +// ============================================================================================ + +__device__ __forceinline__ Fp eval_jit_column( + const JitColumnDesc &desc, + const uint8_t *record, + uint32_t range_max_bits +) { + uint16_t ctype = desc.comp_type; + if ((ctype & JIT_COND_FLAG) != 0) { + uint16_t cond_offset = desc.p[5]; // condition byte offset stored in p[5] + if (record[cond_offset] == 0) return Fp::zero(); + ctype &= ~JIT_COND_FLAG; + } + if ((ctype & JIT_COND_NOT_MAX_U32) != 0) { + uint16_t ptr_offset = desc.p[5]; + uint32_t v; memcpy(&v, record + ptr_offset, 4); + if (v == 0xFFFFFFFFu) return Fp::zero(); + ctype &= ~JIT_COND_NOT_MAX_U32; + } + + switch (ctype) { + case JIT_DIRECT_U32: { + uint32_t v; memcpy(&v, record + desc.p[0], 4); + return Fp(v); + } + case JIT_DIRECT_U8: + return Fp((uint32_t)record[desc.p[0]]); + case JIT_DIRECT_U16: { + uint16_t v; memcpy(&v, record + desc.p[0], 2); + return Fp((uint32_t)v); + } + case JIT_TIMESTAMP_DECOMP: { + uint32_t curr, prev; + memcpy(&curr, record + desc.p[0], 4); + memcpy(&prev, record + desc.p[2], 4); + uint32_t delta = desc.p[1]; + uint32_t diff = (curr + delta) - prev - 1; + uint32_t mask = (1u << range_max_bits) - 1; + uint32_t limb_idx = desc.p[3]; + return Fp((diff >> (range_max_bits * limb_idx)) & mask); + } + case JIT_ALU_RESULT: { + uint8_t opcode = record[desc.p[0]]; + const uint8_t *b = record + desc.p[1]; + const uint8_t *c = record + desc.p[2]; + uint8_t a[4]; run_alu(opcode, b, c, a); + return Fp((uint32_t)a[desc.p[3]]); + } + case JIT_BOOL_FROM_OPCODE: + return Fp(record[desc.p[0]] == (uint8_t)desc.p[1]); + case JIT_POINTER_LIMB: { + uint32_t val; memcpy(&val, record + desc.p[0], 4); + uint16_t imm; memcpy(&imm, record + desc.p[1], 2); + uint8_t sign = record[desc.p[2]]; + uint32_t imm_ext = (uint32_t)imm + (sign ? 0xFFFF0000u : 0u); + uint32_t ptr = val + imm_ext; + return Fp((ptr >> (16 * desc.p[3])) & 0xFFFF); + } + case JIT_SHIFT_RESULT: { + uint8_t opcode = record[desc.p[0]]; + const uint8_t *b = record + desc.p[1]; + const uint8_t *c = record + desc.p[2]; + uint8_t a[4]; run_shift(opcode, b, c, a); + return Fp((uint32_t)a[desc.p[3]]); + } + case JIT_SHIFT_BIT_MUL_L: { + bool is_sll = record[desc.p[0]] == 0; + int ls, bs; get_shift_amounts(record[desc.p[1]], ls, bs); + return is_sll ? Fp(1u << bs) : Fp::zero(); + } + case JIT_SHIFT_BIT_MUL_R: { + bool is_sll = record[desc.p[0]] == 0; + int ls, bs; get_shift_amounts(record[desc.p[1]], ls, bs); + return is_sll ? Fp::zero() : Fp(1u << bs); + } + case JIT_SHIFT_B_SIGN: { + bool is_sra = record[desc.p[0]] == 2; + return is_sra ? Fp((uint32_t)(record[desc.p[1] + 3] >> 7)) : Fp::zero(); + } + case JIT_SHIFT_BIT_MARKER: { + int ls, bs; get_shift_amounts(record[desc.p[0]], ls, bs); + return Fp(bs == (int)desc.p[1]); + } + case JIT_SHIFT_LIMB_MARKER: { + int ls, bs; get_shift_amounts(record[desc.p[0]], ls, bs); + return Fp(ls == (int)desc.p[1]); + } + case JIT_SHIFT_BIT_CARRY: { + bool is_sll = record[desc.p[0]] == 0; + int ls, bs; get_shift_amounts(record[desc.p[2]], ls, bs); + if (bs == 0) return Fp::zero(); + uint8_t bv = record[desc.p[1] + desc.p[3]]; + uint8_t carry = is_sll ? (bv >> (8 - bs)) : (bv & ((1u << bs) - 1)); + return Fp((uint32_t)carry); + } + case JIT_BRANCH_EQ_CMP: { + const uint8_t *a = record + desc.p[0]; + const uint8_t *b = record + desc.p[1]; + bool is_beq = record[desc.p[2]] == 0; + bool eq = (a[0]==b[0] && a[1]==b[1] && a[2]==b[2] && a[3]==b[3]); + return Fp(is_beq ? eq : !eq); + } + case JIT_BRANCH_EQ_DIFF_INV: { + const uint8_t *a = record + desc.p[0]; + const uint8_t *b = record + desc.p[1]; + int diff_idx = 4; + for (int i=0;i<4;i++) { if(a[i]!=b[i]) { diff_idx=i; break; } } + if (diff_idx == 4) diff_idx = 0; + int marker = desc.p[3]; + bool all_eq = (a[0]==b[0] && a[1]==b[1] && a[2]==b[2] && a[3]==b[3]); + if (marker == diff_idx && !all_eq) { + Fp av = Fp((uint32_t)a[diff_idx]); + Fp bv = Fp((uint32_t)b[diff_idx]); + return inv(av - bv); + } + return Fp::zero(); + } + case JIT_CONSTANT: + return Fp(desc.p[0] | ((uint32_t)desc.p[1] << 16)); + case JIT_LS_RD_RS2_PTR: { + uint32_t v; memcpy(&v, record + desc.p[0], 4); + return (v == 0xFFFFFFFFu) ? Fp::zero() : Fp(v); + } + case JIT_LS_NEEDS_WRITE: { + uint32_t v; memcpy(&v, record + desc.p[0], 4); + return Fp(v != 0xFFFFFFFFu); + } + case JIT_LS_WRITE_AUX_TS: { + uint32_t rd; memcpy(&rd, record + desc.p[1], 4); + if (rd == 0xFFFFFFFFu) return Fp::zero(); + uint32_t v; memcpy(&v, record + desc.p[0], 4); + return Fp(v); + } + case JIT_LS_WRITE_AUX_DEC: { + uint32_t rd; memcpy(&rd, record + desc.p[2], 4); + if (rd == 0xFFFFFFFFu) return Fp::zero(); + uint32_t curr, prev; + memcpy(&curr, record + desc.p[0], 4); + memcpy(&prev, record + desc.p[1], 4); + uint32_t diff = (curr + 2) - prev - 1; + uint32_t mask = (1u << range_max_bits) - 1; + return Fp((diff >> (range_max_bits * desc.p[3])) & mask); + } + case JIT_LS_IS_LOAD: + return Fp(record[desc.p[0]] <= 2); + case JIT_LS_FLAG: { + uint8_t op = record[desc.p[0]]; + uint8_t sh = record[desc.p[1]]; + uint8_t flags[4] = {0,0,0,0}; + switch(op) { + case 0: flags[0]=2; break; + case 2: if(sh==0) flags[1]=2; else if(sh==2) flags[2]=2; break; + case 1: switch(sh){case 0:flags[3]=2;break;case 1:flags[0]=1;break;case 2:flags[1]=1;break;case 3:flags[2]=1;break;} break; + case 3: flags[3]=1; break; + case 4: if(sh==0){flags[0]=1;flags[1]=1;}else if(sh==2){flags[0]=1;flags[2]=1;} break; + case 5: switch(sh){case 0:flags[0]=1;flags[3]=1;break;case 1:flags[1]=1;flags[2]=1;break;case 2:flags[1]=1;flags[3]=1;break;case 3:flags[2]=1;flags[3]=1;break;} break; + } + return Fp((uint32_t)flags[desc.p[2]]); + } + case JIT_LS_WRITE_DATA: { + uint8_t op = record[desc.p[0]]; + int sh = (int)record[desc.p[1]]; + const uint8_t *rd = record + desc.p[2]; + int idx = desc.p[4]; + uint32_t prev; memcpy(&prev, record + desc.p[3] + idx*4, 4); + uint32_t wd = 0; + switch(op) { + case 0: case 3: wd = rd[idx]; break; // LOADW, STOREW + case 2: wd = (idx < 2) ? rd[idx+sh] : 0; break; // LOADHU + case 1: wd = (idx == 0) ? rd[sh] : 0; break; // LOADBU + case 4: wd = (idx>=sh && idx<2+sh) ? rd[idx-sh] : prev; break; // STOREH + case 5: wd = (idx==sh) ? rd[0] : prev; break; // STOREB + } + return Fp(wd); + } + // ── LessThan arms (p[0]=opcode_off, p[1]=b_off, p[2]=c_off, p[3]=marker_index) ── + case JIT_LT_CMP: { + bool is_slt = (record[desc.p[0]] == 0); + int diff_idx; bool b_sign, c_sign; + bool cmp = run_less_than(is_slt, record + desc.p[1], record + desc.p[2], + &diff_idx, &b_sign, &c_sign); + return Fp((uint32_t)cmp); + } + case JIT_LT_DIFF_VAL: { + bool is_slt = (record[desc.p[0]] == 0); + const uint8_t *b = record + desc.p[1]; + const uint8_t *c = record + desc.p[2]; + int diff_idx; bool b_sign, c_sign; + bool cmp = run_less_than(is_slt, b, c, &diff_idx, &b_sign, &c_sign); + if (diff_idx == 4) return Fp::zero(); + if (diff_idx == 3) { + // Use signed-aware MSb encodings — match the field arithmetic. + Fp b_msb = b_sign ? -Fp((uint32_t)(256u - b[3])) : Fp((uint32_t)b[3]); + Fp c_msb = c_sign ? -Fp((uint32_t)(256u - c[3])) : Fp((uint32_t)c[3]); + return cmp ? (c_msb - b_msb) : (b_msb - c_msb); + } + uint8_t big = cmp ? c[diff_idx] : b[diff_idx]; + uint8_t small = cmp ? b[diff_idx] : c[diff_idx]; + return Fp((uint32_t)(big - small)); + } + case JIT_LT_DIFF_MARKER: { + bool is_slt = (record[desc.p[0]] == 0); + int diff_idx; bool b_sign, c_sign; + run_less_than(is_slt, record + desc.p[1], record + desc.p[2], + &diff_idx, &b_sign, &c_sign); + int marker = (int)desc.p[3]; + return Fp((uint32_t)((diff_idx != 4) && (diff_idx == marker))); + } + case JIT_LT_B_MSB_F: { + bool is_slt = (record[desc.p[0]] == 0); + uint8_t b3 = record[desc.p[1] + 3]; + bool b_sign = is_slt && ((b3 >> 7) == 1); + return b_sign ? -Fp((uint32_t)(256u - b3)) : Fp((uint32_t)b3); + } + case JIT_LT_C_MSB_F: { + bool is_slt = (record[desc.p[0]] == 0); + uint8_t c3 = record[desc.p[1] + 3]; + bool c_sign = is_slt && ((c3 >> 7) == 1); + return c_sign ? -Fp((uint32_t)(256u - c3)) : Fp((uint32_t)c3); + } + // ── BranchLessThan arms (p[0]=opcode_off, p[1]=a_off, p[2]=b_off) ── + case JIT_BLT_CMP: { + uint8_t op = record[desc.p[0]]; + int diff_idx; bool a_sign, b_sign; + bool cmp = run_branch_lt(op, record + desc.p[1], record + desc.p[2], + &diff_idx, &a_sign, &b_sign); + return Fp((uint32_t)cmp); + } + case JIT_BLT_CMP_LT: { + uint8_t op = record[desc.p[0]]; + int diff_idx; bool a_sign, b_sign; + bool cmp = run_branch_lt(op, record + desc.p[1], record + desc.p[2], + &diff_idx, &a_sign, &b_sign); + bool ge_op = (op == 2) || (op == 3); + return Fp((uint32_t)(cmp ^ ge_op)); + } + case JIT_BLT_DIFF_VAL: { + uint8_t op = record[desc.p[0]]; + const uint8_t *a = record + desc.p[1]; + const uint8_t *b = record + desc.p[2]; + int diff_idx; bool a_sign, b_sign; + bool cmp = run_branch_lt(op, a, b, &diff_idx, &a_sign, &b_sign); + bool ge_op = (op == 2) || (op == 3); + bool cmp_lt = cmp ^ ge_op; + if (diff_idx == 4) return Fp::zero(); + if (diff_idx == 3) { + Fp a_msb = a_sign ? -Fp((uint32_t)(256u - a[3])) : Fp((uint32_t)a[3]); + Fp b_msb = b_sign ? -Fp((uint32_t)(256u - b[3])) : Fp((uint32_t)b[3]); + return cmp_lt ? (b_msb - a_msb) : (a_msb - b_msb); + } + uint8_t big = cmp_lt ? b[diff_idx] : a[diff_idx]; + uint8_t small = cmp_lt ? a[diff_idx] : b[diff_idx]; + return Fp((uint32_t)(big - small)); + } + case JIT_BLT_DIFF_MARKER: { + uint8_t op = record[desc.p[0]]; + int diff_idx; bool a_sign, b_sign; + run_branch_lt(op, record + desc.p[1], record + desc.p[2], + &diff_idx, &a_sign, &b_sign); + int marker = (int)desc.p[3]; + return Fp((uint32_t)((diff_idx != 4) && (diff_idx == marker))); + } + case JIT_BLT_A_MSB_F: { + uint8_t op = record[desc.p[0]]; + bool signed_op = (op == 0) || (op == 2); + uint8_t a3 = record[desc.p[1] + 3]; + bool a_sign = signed_op && ((a3 >> 7) == 1); + return a_sign ? -Fp((uint32_t)(256u - a3)) : Fp((uint32_t)a3); + } + case JIT_BLT_B_MSB_F: { + uint8_t op = record[desc.p[0]]; + bool signed_op = (op == 0) || (op == 2); + uint8_t b3 = record[desc.p[1] + 3]; + bool b_sign = signed_op && ((b3 >> 7) == 1); + return b_sign ? -Fp((uint32_t)(256u - b3)) : Fp((uint32_t)b3); + } + // ── Auipc rd_data limb ── + // p[0] = pc_byte_offset, p[1] = imm_byte_offset, p[2] = limb_index. + case JIT_AUIPC_RD_LIMB: { + uint32_t pc, imm; + memcpy(&pc, record + desc.p[0], 4); + memcpy(&imm, record + desc.p[1], 4); + uint32_t rd = pc + (imm << 8); + return Fp((uint32_t)((rd >> (8 * desc.p[2])) & 0xFFu)); + } + // ── Jalr to_pc least-significant bit ── + // p[0] = rs1_byte_offset (u32), p[1] = imm_byte_offset (u16), + // p[2] = imm_sign_byte_offset (1 byte: 0 or 1). + case JIT_JALR_TO_PC_LSB: { + uint32_t rs1; memcpy(&rs1, record + desc.p[0], 4); + uint16_t imm; memcpy(&imm, record + desc.p[1], 2); + uint8_t sign = record[desc.p[2]]; + uint32_t to_pc = rs1 + (uint32_t)imm + (sign ? 0xFFFF0000u : 0u); + return Fp((uint32_t)(to_pc & 1u)); + } + // ── Jalr to_pc_limbs[limb_index] ── + // limb 0: (to_pc & 0xFFFF) >> 1 + // limb 1: to_pc >> 16 + case JIT_JALR_TO_PC_LIMB: { + uint32_t rs1; memcpy(&rs1, record + desc.p[0], 4); + uint16_t imm; memcpy(&imm, record + desc.p[1], 2); + uint8_t sign = record[desc.p[2]]; + uint32_t to_pc = rs1 + (uint32_t)imm + (sign ? 0xFFFF0000u : 0u); + uint32_t v = (desc.p[3] == 0u) ? ((to_pc & 0xFFFFu) >> 1) : (to_pc >> 16); + return Fp(v); + } + // ── Jalr rd_data limb (top 3 bytes of pc + 4) ── + // p[0] = pc_byte_offset, p[1] = limb_index ∈ 0..3 + case JIT_JALR_RD_LIMB: { + uint32_t pc; memcpy(&pc, record + desc.p[0], 4); + uint32_t rd = pc + 4u; + return Fp((uint32_t)((rd >> (8u * (desc.p[1] + 1u))) & 0xFFu)); + } + default: + return Fp::zero(); + } +} + +// ============================================================================================ +// Main JIT tracegen kernel +// ============================================================================================ + +__global__ void apc_jit_tracegen_kernel( + Fp* __restrict__ d_output, // APC trace, column-major (pre-zeroed) + size_t H, // trace height (power of 2) + int num_apc_calls, // actual number of valid rows + const uint8_t* __restrict__ d_arena, // concatenated arena buffer + const JitInstructionDesc* __restrict__ d_instructions, + int n_instructions, + const JitColumnDesc* __restrict__ d_col_descs, + uint32_t range_max_bits +) { + const size_t total_threads = (size_t)gridDim.x * (size_t)blockDim.x; + const size_t tid = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; + + for (size_t r = tid; r < (size_t)num_apc_calls; r += total_threads) { + for (int i = 0; i < n_instructions; i++) { + const JitInstructionDesc instr = d_instructions[i]; + const uint8_t* record = d_arena + instr.arena_offset + + instr.record_offset + + r * instr.record_stride; + + for (uint32_t c = 0; c < instr.col_desc_count; c++) { + const JitColumnDesc desc = d_col_descs[instr.col_desc_start + c]; + Fp value = eval_jit_column(desc, record, range_max_bits); + d_output[(size_t)desc.apc_col * H + r] = value; + } + } + } + // Padding rows beyond num_apc_calls stay zero (buffer is pre-zeroed) +} + +// ============================================================================================ +// Host launcher +// ============================================================================================ + +extern "C" int _apc_jit_tracegen( + Fp* d_output, + size_t H, + int num_apc_calls, + const uint8_t* d_arena, + const JitInstructionDesc* d_instructions, + int n_instructions, + const JitColumnDesc* d_col_descs, + uint32_t range_max_bits +) { + if (num_apc_calls == 0) return 0; + const int block_x = 256; + const dim3 block(block_x, 1, 1); + unsigned g = (unsigned)((num_apc_calls + block_x - 1) / block_x); + if (g == 0u) g = 1u; + const dim3 grid(g, 1, 1); + + apc_jit_tracegen_kernel<<>>( + d_output, H, num_apc_calls, + d_arena, d_instructions, n_instructions, + d_col_descs, range_max_bits + ); + return (int)cudaGetLastError(); +} diff --git a/openvm/cuda/src/nvrtc_runtime.cu b/openvm/cuda/src/nvrtc_runtime.cu new file mode 100644 index 0000000000..1dd028e688 --- /dev/null +++ b/openvm/cuda/src/nvrtc_runtime.cu @@ -0,0 +1,512 @@ +// Reusable NVRTC + Driver-API shim for per-APC trace-gen kernels. +// +// All extern "C" entry points return 0 on success and a nonzero CUDA / NVRTC +// error code on failure. Caller-owned buffers are documented per function. + +#include +#include +#include +#include +#include + +namespace { + +bool ensure_primary_context() { + // cuCtxSetCurrent is per-thread, so each rayon worker thread must call + // it (cuInit + cuDevicePrimaryCtxRetain are process-global and idempotent). + // Track init per-thread to avoid the syscall on every launch. + static thread_local bool initialized = false; + if (initialized) return true; + if (cuInit(0) != CUDA_SUCCESS) return false; + CUdevice dev = 0; + if (cuDeviceGet(&dev, 0) != CUDA_SUCCESS) return false; + CUcontext ctx = nullptr; + if (cuDevicePrimaryCtxRetain(&ctx, dev) != CUDA_SUCCESS) return false; + if (cuCtxSetCurrent(ctx) != CUDA_SUCCESS) return false; + initialized = true; + return true; +} + +int device_arch_flag(char* out, size_t out_size) { + CUdevice dev = 0; + if (cuDeviceGet(&dev, 0) != CUDA_SUCCESS) return -1; + int major = 0, minor = 0; + if (cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev) + != CUDA_SUCCESS) return -1; + if (cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev) + != CUDA_SUCCESS) return -1; + std::snprintf(out, out_size, "--gpu-architecture=sm_%d%d", major, minor); + return 0; +} + +} // namespace + +extern "C" { + +// Compile a CUDA C++ source string with NVRTC. On success, returns 0 and +// writes a malloc'd PTX buffer into *ptx_out (caller must free with +// powdr_nvrtc_free) and its byte size into *ptx_size_out. On failure, returns +// the (negated) NVRTC result code; if a compilation log is available it is +// written to *log_out (caller must free with powdr_nvrtc_free). +int powdr_nvrtc_compile( + const char* src, + const char* src_name, + char** ptx_out, + size_t* ptx_size_out, + char** log_out +) { + if (!ensure_primary_context()) return -1; + if (ptx_out) *ptx_out = nullptr; + if (ptx_size_out) *ptx_size_out = 0; + if (log_out) *log_out = nullptr; + + char arch_flag[64]; + if (device_arch_flag(arch_flag, sizeof(arch_flag)) != 0) return -2; + + nvrtcProgram prog = nullptr; + nvrtcResult cr = nvrtcCreateProgram(&prog, src, src_name, 0, nullptr, nullptr); + if (cr != NVRTC_SUCCESS) return -1000 - (int)cr; + + const char* opts[] = { + arch_flag, + "--use_fast_math", + "-default-device", + "--std=c++17", + }; + cr = nvrtcCompileProgram(prog, (int)(sizeof(opts) / sizeof(opts[0])), opts); + + // Fetch the log either way so caller can debug compile errors. + size_t log_size = 0; + nvrtcGetProgramLogSize(prog, &log_size); + if (log_size > 1 && log_out) { + char* log = (char*)std::malloc(log_size); + if (log != nullptr) { + if (nvrtcGetProgramLog(prog, log) == NVRTC_SUCCESS) { + *log_out = log; + } else { + std::free(log); + } + } + } + + if (cr != NVRTC_SUCCESS) { + nvrtcDestroyProgram(&prog); + return -1000 - (int)cr; + } + + size_t ptx_size = 0; + cr = nvrtcGetPTXSize(prog, &ptx_size); + if (cr != NVRTC_SUCCESS) { nvrtcDestroyProgram(&prog); return -1000 - (int)cr; } + char* ptx = (char*)std::malloc(ptx_size); + if (ptx == nullptr) { nvrtcDestroyProgram(&prog); return -3; } + cr = nvrtcGetPTX(prog, ptx); + if (cr != NVRTC_SUCCESS) { + std::free(ptx); + nvrtcDestroyProgram(&prog); + return -1000 - (int)cr; + } + nvrtcDestroyProgram(&prog); + + if (ptx_out) *ptx_out = ptx; else std::free(ptx); + if (ptx_size_out) *ptx_size_out = ptx_size; + return 0; +} + +// Free a buffer allocated by powdr_nvrtc_compile (PTX or log). +void powdr_nvrtc_free(char* p) { std::free(p); } + +// Load a PTX blob into a CUDA module. *module_out holds an opaque CUmodule +// handle on success (cast through void*). Caller must call +// powdr_nvrtc_unload_module to release it. +int powdr_nvrtc_load_module(const void* ptx, size_t /*ptx_size*/, void** module_out) { + if (!ensure_primary_context()) return -1; + if (module_out) *module_out = nullptr; + CUmodule mod = nullptr; + CUresult r = cuModuleLoadData(&mod, ptx); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + if (module_out) *module_out = (void*)mod; + return 0; +} + +int powdr_nvrtc_get_function(void* module, const char* name, void** fn_out) { + if (fn_out) *fn_out = nullptr; + CUfunction fn = nullptr; + CUresult r = cuModuleGetFunction(&fn, (CUmodule)module, name); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + if (fn_out) *fn_out = (void*)fn; + return 0; +} + +int powdr_nvrtc_unload_module(void* module) { + CUresult r = cuModuleUnload((CUmodule)module); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + return 0; +} + +// Launch a JIT-compiled trace-gen kernel that conforms to the v1 signature: +// +// __global__ void (uint32_t* d_output, +// size_t H, +// int N, +// const uint8_t* d_arena, +// uint32_t range_max_bits); +// +// Synchronizes before returning so any kernel error surfaces to the caller. +int powdr_nvrtc_launch_jit_v1( + void* fn, + unsigned int* d_output, + size_t H, + int num_apc_calls, + const unsigned char* d_arena, + unsigned int range_max_bits, + unsigned int grid_x, + unsigned int block_x +) { + if (!ensure_primary_context()) return -1; + if (fn == nullptr) return -2; + if (grid_x == 0) grid_x = 1; + if (block_x == 0) block_x = 256; + + void* args[] = { + (void*)&d_output, + (void*)&H, + (void*)&num_apc_calls, + (void*)&d_arena, + (void*)&range_max_bits, + }; + + CUresult r = cuLaunchKernel( + (CUfunction)fn, + grid_x, 1, 1, + block_x, 1, 1, + 0, + nullptr, + args, + nullptr); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + + r = cuCtxSynchronize(); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + return 0; +} + +// Upload `host_data` (size_bytes) into a `__constant__` symbol on the +// kernel's module. Caller passes the module from cuModuleLoadData and the +// symbol name as it appears in CUDA source. Returns 0 on success. +// +// Used by the kind-templated bus emitter to push per-APC op tables into +// the kernel's `__constant__` array between launches. +int powdr_nvrtc_set_constant_symbol( + void* module, + const char* sym_name, + const void* host_data, + size_t size_bytes +) { + if (!ensure_primary_context()) return -1; + if (module == nullptr || sym_name == nullptr) return -2; + if (size_bytes == 0) return 0; + + CUdeviceptr d_ptr = 0; + size_t d_size = 0; + CUresult r = cuModuleGetGlobal(&d_ptr, &d_size, (CUmodule)module, sym_name); + if (r != CUDA_SUCCESS) return -3000 - (int)r; + if (size_bytes > d_size) return -2; // overflow against declared array + + r = cuMemcpyHtoD(d_ptr, host_data, size_bytes); + if (r != CUDA_SUCCESS) return -3000 - (int)r; + return 0; +} + +// Launch a per-APC codegen bus kernel matching the bus_v4 signature: no +// d_ops pointer (constants are baked into the kernel source) and no +// n_ops arg (constexpr inside the kernel; grid_x is sized at host time). +// +// __global__ void ( +// const unsigned int* d_output, int N, unsigned long long H, +// unsigned int* d_hist, +// ); +// +// `extra1` is used only by tuple2; pass 0 + has_extra1=0 for others. +int powdr_nvrtc_launch_bus_v4( + void* fn, + const unsigned int* d_output, + int num_apc_calls, + unsigned long long H, + unsigned int* d_hist, + unsigned int extra0, + unsigned int extra1, + unsigned int has_extra1, + unsigned int grid_x, + unsigned int block_x +) { + if (!ensure_primary_context()) return -1; + if (fn == nullptr) return -2; + if (grid_x == 0) grid_x = 1; + if (block_x == 0) block_x = 256; + + void* args_no_extra1[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&d_hist, + (void*)&extra0, + }; + void* args_with_extra1[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&d_hist, + (void*)&extra0, + (void*)&extra1, + }; + // bitwise kernels have NO extra0 in v4 either; the extra0 field is + // only used by var_range (var_num_bins) and tuple2 (sz0). Caller is + // responsible for matching the signature. + + CUresult r = cuLaunchKernel( + (CUfunction)fn, + grid_x, 1, 1, + block_x, 1, 1, + 0, + nullptr, + has_extra1 ? args_with_extra1 : args_no_extra1, + nullptr); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + + r = cuCtxSynchronize(); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + return 0; +} + +// Launch a per-APC codegen bus kernel for bitwise (no extra0). This is a +// thin variant of v4 that omits the extra0 placeholder. +int powdr_nvrtc_launch_bus_v4_bitwise( + void* fn, + const unsigned int* d_output, + int num_apc_calls, + unsigned long long H, + unsigned int* d_hist, + unsigned int grid_x, + unsigned int block_x +) { + if (!ensure_primary_context()) return -1; + if (fn == nullptr) return -2; + if (grid_x == 0) grid_x = 1; + if (block_x == 0) block_x = 256; + + void* args[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&d_hist, + }; + + CUresult r = cuLaunchKernel( + (CUfunction)fn, + grid_x, 1, 1, + block_x, 1, 1, + 0, + nullptr, + args, + nullptr); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + + r = cuCtxSynchronize(); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + return 0; +} + +// Launch a kind-templated bus kernel matching the bus_v3 signature: same +// shape as v2 but takes a `d_ops` pointer instead of relying on a +// __constant__ symbol upload (op tables now live in global memory so they +// can grow past 64KB per module). +// +// __global__ void ( +// const unsigned int* d_output, int N, unsigned long long H, +// unsigned int* d_hist, +// , +// unsigned int n_ops, +// const void* d_ops); +// +// `extra1` is used only by tuple2 (second size arg); pass 0 + has_extra1=0 +// for var_range and bitwise. Synchronizes before returning. +int powdr_nvrtc_launch_bus_v3( + void* fn, + const unsigned int* d_output, + int num_apc_calls, + unsigned long long H, + unsigned int* d_hist, + unsigned int extra0, + unsigned int extra1, + unsigned int has_extra1, + unsigned int n_ops, + const void* d_ops, + unsigned int grid_x, + unsigned int block_x +) { + if (!ensure_primary_context()) return -1; + if (fn == nullptr) return -2; + if (grid_x == 0) grid_x = 1; + if (block_x == 0) block_x = 256; + + void* args_no_extra1[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&d_hist, + (void*)&extra0, + (void*)&n_ops, + (void*)&d_ops, + }; + void* args_with_extra1[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&d_hist, + (void*)&extra0, + (void*)&extra1, + (void*)&n_ops, + (void*)&d_ops, + }; + + CUresult r = cuLaunchKernel( + (CUfunction)fn, + grid_x, 1, 1, + block_x, 1, 1, + 0, + nullptr, + has_extra1 ? args_with_extra1 : args_no_extra1, + nullptr); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + + r = cuCtxSynchronize(); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + return 0; +} + +// Launch a kind-templated bus kernel matching the bus_v2 signature shared +// by var_range / tuple2 / bitwise_range / bitwise_xor: +// +// __global__ void ( +// const unsigned int* d_output, int N, unsigned long long H, +// unsigned int* d_hist, +// , +// unsigned int n_ops); +// +// Synchronizes before returning. The histogram-specific extra args are +// passed via `extra0` (always present) and `extra1` (used only by tuple2 +// for the second size; pass 0 for the others). +int powdr_nvrtc_launch_bus_v2( + void* fn, + const unsigned int* d_output, + int num_apc_calls, + unsigned long long H, + unsigned int* d_hist, + unsigned int extra0, + unsigned int extra1, + unsigned int has_extra1, + unsigned int n_ops, + unsigned int grid_x, + unsigned int block_x +) { + if (!ensure_primary_context()) return -1; + if (fn == nullptr) return -2; + if (grid_x == 0) grid_x = 1; + if (block_x == 0) block_x = 256; + + void* args_no_extra1[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&d_hist, + (void*)&extra0, + (void*)&n_ops, + }; + void* args_with_extra1[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&d_hist, + (void*)&extra0, + (void*)&extra1, + (void*)&n_ops, + }; + + CUresult r = cuLaunchKernel( + (CUfunction)fn, + grid_x, 1, 1, + block_x, 1, 1, + 0, + nullptr, + has_extra1 ? args_with_extra1 : args_no_extra1, + nullptr); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + + r = cuCtxSynchronize(); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + return 0; +} + +// Launch a JIT-compiled bus kernel that conforms to the bus_v1 signature: +// +// __global__ void ( +// const unsigned int* d_output, int N, unsigned long long H, +// unsigned int var_range_bus_id, unsigned int* d_var_hist, unsigned int var_num_bins, +// unsigned int tuple2_bus_id, unsigned int* d_tuple2_hist, +// unsigned int tuple2_sz0, unsigned int tuple2_sz1, +// unsigned int bitwise_bus_id, unsigned int* d_bitwise_hist); +// +// Synchronizes before returning so any kernel error surfaces to the caller. +int powdr_nvrtc_launch_bus_v1( + void* fn, + const unsigned int* d_output, + int num_apc_calls, + unsigned long long H, + unsigned int var_range_bus_id, + unsigned int* d_var_hist, + unsigned int var_num_bins, + unsigned int tuple2_bus_id, + unsigned int* d_tuple2_hist, + unsigned int tuple2_sz0, + unsigned int tuple2_sz1, + unsigned int bitwise_bus_id, + unsigned int* d_bitwise_hist, + unsigned int grid_x, + unsigned int block_x +) { + if (!ensure_primary_context()) return -1; + if (fn == nullptr) return -2; + if (grid_x == 0) grid_x = 1; + if (block_x == 0) block_x = 256; + + void* args[] = { + (void*)&d_output, + (void*)&num_apc_calls, + (void*)&H, + (void*)&var_range_bus_id, + (void*)&d_var_hist, + (void*)&var_num_bins, + (void*)&tuple2_bus_id, + (void*)&d_tuple2_hist, + (void*)&tuple2_sz0, + (void*)&tuple2_sz1, + (void*)&bitwise_bus_id, + (void*)&d_bitwise_hist, + }; + + CUresult r = cuLaunchKernel( + (CUfunction)fn, + grid_x, 1, 1, + block_x, 1, 1, + 0, + nullptr, + args, + nullptr); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + + r = cuCtxSynchronize(); + if (r != CUDA_SUCCESS) return -2000 - (int)r; + return 0; +} + +} // extern "C" diff --git a/openvm/cuda/src/nvrtc_spike.cu b/openvm/cuda/src/nvrtc_spike.cu new file mode 100644 index 0000000000..1157f4fd03 --- /dev/null +++ b/openvm/cuda/src/nvrtc_spike.cu @@ -0,0 +1,120 @@ +// Phase 0 spike: prove that NVRTC-compiled kernels can launch against +// DeviceBuffer-allocated memory (cudaMallocAsync / VPMM) using the Driver API. +// +// Exposes one extern "C" function: +// powdr_nvrtc_spike_run_noop(d_ptr, n) +// which compiles a trivial kernel that writes 0x12345678 into the first `n` +// elements of `d_ptr`, loads the resulting PTX as a module, launches the +// kernel, and synchronizes. Returns 0 on success; nonzero on any error. + +#include +#include +#include +#include + +namespace { + +const char* k_noop_src = R"CUDA_SRC( +extern "C" __global__ void noop(unsigned int* p, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) p[i] = 0x12345678u; +} +)CUDA_SRC"; + +#define NVRTC_CHECK(call) \ + do { \ + nvrtcResult _r = (call); \ + if (_r != NVRTC_SUCCESS) { \ + std::fprintf(stderr, "NVRTC error %d at %s:%d (%s)\n", \ + (int)_r, __FILE__, __LINE__, nvrtcGetErrorString(_r)); \ + return -((int)_r + 1000); \ + } \ + } while (0) + +#define CU_CHECK(call) \ + do { \ + CUresult _r = (call); \ + if (_r != CUDA_SUCCESS) { \ + const char* _msg = nullptr; \ + cuGetErrorString(_r, &_msg); \ + std::fprintf(stderr, "CUDA driver error %d at %s:%d (%s)\n", \ + (int)_r, __FILE__, __LINE__, \ + _msg ? _msg : "unknown"); \ + return -((int)_r + 2000); \ + } \ + } while (0) + +} + +extern "C" int powdr_nvrtc_spike_run_noop(unsigned int* d_ptr, int n) { + if (d_ptr == nullptr || n <= 0) return -1; + + // Ensure a CUDA context exists. cuInit + cuDevicePrimaryCtxRetain attaches + // to the same primary context the Runtime API uses, so the device pointer + // allocated by cudaMallocAsync / VPMM is valid here. + CU_CHECK(cuInit(0)); + CUdevice dev = 0; + CU_CHECK(cuDeviceGet(&dev, 0)); + CUcontext ctx = nullptr; + CU_CHECK(cuDevicePrimaryCtxRetain(&ctx, dev)); + CU_CHECK(cuCtxSetCurrent(ctx)); + + // Determine arch for nvrtc: compute capability of device 0. + int major = 0, minor = 0; + CU_CHECK(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev)); + CU_CHECK(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev)); + char arch_flag[64]; + std::snprintf(arch_flag, sizeof(arch_flag), "--gpu-architecture=sm_%d%d", major, minor); + + nvrtcProgram prog = nullptr; + NVRTC_CHECK(nvrtcCreateProgram(&prog, k_noop_src, "noop.cu", 0, nullptr, nullptr)); + + const char* opts[] = { + arch_flag, + "--use_fast_math", + "-default-device", + "--std=c++17", + }; + nvrtcResult comp_r = nvrtcCompileProgram(prog, (int)(sizeof(opts) / sizeof(opts[0])), opts); + if (comp_r != NVRTC_SUCCESS) { + size_t log_size = 0; + nvrtcGetProgramLogSize(prog, &log_size); + if (log_size > 1) { + char* log = (char*)std::malloc(log_size); + nvrtcGetProgramLog(prog, log); + std::fprintf(stderr, "NVRTC compile log:\n%s\n", log); + std::free(log); + } + nvrtcDestroyProgram(&prog); + return -((int)comp_r + 1000); + } + + size_t ptx_size = 0; + NVRTC_CHECK(nvrtcGetPTXSize(prog, &ptx_size)); + char* ptx = (char*)std::malloc(ptx_size); + NVRTC_CHECK(nvrtcGetPTX(prog, ptx)); + NVRTC_CHECK(nvrtcDestroyProgram(&prog)); + + CUmodule mod = nullptr; + CU_CHECK(cuModuleLoadData(&mod, ptx)); + std::free(ptx); + + CUfunction fn = nullptr; + CU_CHECK(cuModuleGetFunction(&fn, mod, "noop")); + + int block_x = 256; + int grid_x = (n + block_x - 1) / block_x; + + void* args[] = { (void*)&d_ptr, (void*)&n }; + CU_CHECK(cuLaunchKernel(fn, + grid_x, 1, 1, + block_x, 1, 1, + 0, + nullptr, // default stream (driver API => null = legacy default stream) + args, + nullptr)); + + CU_CHECK(cuCtxSynchronize()); + CU_CHECK(cuModuleUnload(mod)); + return 0; +} diff --git a/openvm/src/air_builder.rs b/openvm/src/air_builder.rs index 5e4c44cb2b..d3d323a2fb 100644 --- a/openvm/src/air_builder.rs +++ b/openvm/src/air_builder.rs @@ -2,87 +2,33 @@ use std::sync::Arc; use openvm_stark_backend::air_builders::symbolic::get_symbolic_builder; use openvm_stark_backend::air_builders::symbolic::SymbolicRapBuilder; -use openvm_stark_backend::config::Com; -use openvm_stark_backend::config::StarkGenericConfig; -use openvm_stark_backend::config::Val; -use openvm_stark_backend::interaction::RapPhaseSeqKind; -use openvm_stark_backend::keygen::types::ProverOnlySinglePreprocessedData; use openvm_stark_backend::keygen::types::TraceWidth; -use openvm_stark_backend::keygen::types::VerifierSinglePreprocessedData; -use openvm_stark_backend::p3_commit::Pcs; -use openvm_stark_backend::p3_matrix::Matrix; -use openvm_stark_backend::rap::AnyRap; +use openvm_stark_backend::p3_air::BaseAir; +use openvm_stark_backend::AnyAir; +use openvm_stark_backend::StarkProtocolConfig; +use openvm_stark_backend::Val; -pub struct PrepKeygenData { - pub _verifier_data: Option>>, - pub prover_data: Option>, +pub struct AirKeygenBuilder { + air: Arc>, } -pub struct AirKeygenBuilder { - air: Arc>, - prep_keygen_data: PrepKeygenData, -} - -fn compute_prep_data_for_air( - pcs: &SC::Pcs, - air: &dyn AnyRap, -) -> PrepKeygenData { - let preprocessed_trace = air.preprocessed_trace(); - let vpdata_opt = preprocessed_trace.map(|trace| { - let domain = pcs.natural_domain_for_degree(trace.height()); - let (commit, data) = pcs.commit(vec![(domain, trace.clone())]); - let vdata = VerifierSinglePreprocessedData { commit }; - let pdata = ProverOnlySinglePreprocessedData { - trace: Arc::new(trace), - data: Arc::new(data), - }; - (vdata, pdata) - }); - if let Some((vdata, pdata)) = vpdata_opt { - PrepKeygenData { - prover_data: Some(pdata), - _verifier_data: Some(vdata), - } - } else { - PrepKeygenData { - prover_data: None, - _verifier_data: None, - } - } -} - -impl AirKeygenBuilder { - pub fn new(pcs: &SC::Pcs, air: Arc>) -> Self { - let prep_keygen_data = compute_prep_data_for_air(pcs, air.as_ref()); - AirKeygenBuilder { - air, - prep_keygen_data, - } +impl AirKeygenBuilder { + pub fn new(air: Arc>) -> Self { + AirKeygenBuilder { air } } pub fn get_symbolic_builder( &self, - max_constraint_degree: Option, + _max_constraint_degree: Option, ) -> SymbolicRapBuilder> { + let preprocessed_width: Option = + BaseAir::>::preprocessed_trace(self.air.as_ref()).map(|t| t.width); let width = TraceWidth { - preprocessed: self.prep_keygen_data.width(), + preprocessed: preprocessed_width, cached_mains: self.air.cached_main_widths(), common_main: self.air.common_main_width(), after_challenge: vec![], }; - get_symbolic_builder( - self.air.as_ref(), - &width, - &[], - &[], - RapPhaseSeqKind::None, - max_constraint_degree.unwrap_or(0), - ) - } -} - -impl PrepKeygenData { - pub fn width(&self) -> Option { - self.prover_data.as_ref().map(|d| d.trace.width()) + get_symbolic_builder(self.air.as_ref(), &width, &[], &[]) } } diff --git a/openvm/src/bin/nvrtc_spike.rs b/openvm/src/bin/nvrtc_spike.rs new file mode 100644 index 0000000000..01cbb86762 --- /dev/null +++ b/openvm/src/bin/nvrtc_spike.rs @@ -0,0 +1,50 @@ +// Phase 0 NVRTC spike: confirm an NVRTC-compiled kernel can launch against +// memory allocated by openvm-cuda-common's DeviceBuffer (cudaMallocAsync / +// VPMM) and that we can read back correct values from the host. +// +// Run with: +// cargo run -p powdr-openvm --features cuda --bin nvrtc_spike --release + +#[cfg(not(feature = "cuda"))] +fn main() { + eprintln!("nvrtc_spike requires the `cuda` feature. Re-run with --features cuda."); + std::process::exit(2); +} + +#[cfg(feature = "cuda")] +fn main() { + use openvm_cuda_common::{copy::MemCopyD2H, d_buffer::DeviceBuffer}; + use powdr_openvm::cuda_abi::powdr_nvrtc_spike_run_noop; + + const N: usize = 4096; + + let d_buf: DeviceBuffer = DeviceBuffer::with_capacity(N); + d_buf.fill_zero().expect("fill_zero failed"); + + let rc = unsafe { powdr_nvrtc_spike_run_noop(d_buf.as_mut_ptr(), N as i32) }; + assert_eq!( + rc, 0, + "powdr_nvrtc_spike_run_noop failed with code {} (see stderr for NVRTC/CUDA error)", + rc + ); + + let host: Vec = d_buf.to_host().expect("D2H copy failed"); + assert_eq!(host.len(), N, "host vec length mismatch"); + let bad: Vec<(usize, u32)> = host + .iter() + .enumerate() + .filter(|(_, v)| **v != 0x1234_5678u32) + .take(8) + .map(|(i, v)| (i, *v)) + .collect(); + assert!( + bad.is_empty(), + "NVRTC kernel produced wrong values; first mismatches: {:?}", + bad + ); + + println!( + "Phase 0 spike OK — NVRTC-compiled kernel wrote 0x12345678 into {} cells of a DeviceBuffer.", + N + ); +} diff --git a/openvm/src/cuda_abi.rs b/openvm/src/cuda_abi.rs index b0f0a67ae6..d68be9682f 100644 --- a/openvm/src/cuda_abi.rs +++ b/openvm/src/cuda_abi.rs @@ -2,7 +2,7 @@ use openvm_cuda_backend::base::DeviceMatrix; use openvm_cuda_common::{d_buffer::DeviceBuffer, error::CudaError}; -use openvm_stark_backend::prover::hal::MatrixDimensions; +use openvm_stark_backend::prover::MatrixDimensions; use openvm_stark_sdk::p3_baby_bear::BabyBear; extern "C" { @@ -61,6 +61,204 @@ extern "C" { bitwise_bus_id: u32, // bus id for the bitwise lookup d_bitwise_hist: *mut u32, // device histogram for bitwise lookup ) -> i32; + + /// Phase 0 spike: same dispatch + atomicAdd shape as `_apc_apply_bus` + /// but with bytecode evaluation removed (one trace-cell load + one + /// histogram update per row per interaction). Used to determine whether + /// the bytecode VM is the bottleneck (NVRTC will help) or atomic-add + /// traffic is (NVRTC won't help). + pub fn _apc_apply_bus_spike( + d_output: *const BabyBear, + num_apc_calls: i32, + d_interactions: *const DevInteraction, + n_interactions: usize, + var_range_bus_id: u32, + d_var_hist: *mut u32, + var_num_bins: usize, + 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; + + /// JIT trace generation: reads record bytes directly from arenas, + /// computes only surviving APC columns via descriptor arrays. + pub fn _apc_jit_tracegen( + d_output: *mut BabyBear, + output_height: usize, + num_apc_calls: i32, + d_arena: *const u8, // concatenated arena bytes + d_instructions: *const JitInstructionDesc, // instruction descriptors + n_instructions: i32, + d_col_descs: *const JitColumnDesc, // column descriptors + range_max_bits: u32, + ) -> i32; + + /// Phase 0 NVRTC spike: compile a trivial kernel at runtime, launch it + /// against the supplied device pointer, write 0x12345678 into n elements. + /// Returns 0 on success, nonzero on any CUDA / NVRTC error. + pub fn powdr_nvrtc_spike_run_noop(d_ptr: *mut u32, n: i32) -> i32; + + /// Compile a CUDA source string with NVRTC. On success, *ptx_out points + /// to a malloc'd PTX buffer of *ptx_size_out bytes; caller must free + /// with powdr_nvrtc_free. On failure (return != 0), *log_out may point + /// to a malloc'd diagnostics string. + pub fn powdr_nvrtc_compile( + src: *const std::ffi::c_char, + src_name: *const std::ffi::c_char, + ptx_out: *mut *mut std::ffi::c_char, + ptx_size_out: *mut usize, + log_out: *mut *mut std::ffi::c_char, + ) -> i32; + + pub fn powdr_nvrtc_free(p: *mut std::ffi::c_char); + + pub fn powdr_nvrtc_load_module( + ptx: *const std::ffi::c_void, + ptx_size: usize, + module_out: *mut *mut std::ffi::c_void, + ) -> i32; + + pub fn powdr_nvrtc_get_function( + module: *mut std::ffi::c_void, + name: *const std::ffi::c_char, + fn_out: *mut *mut std::ffi::c_void, + ) -> i32; + + pub fn powdr_nvrtc_unload_module(module: *mut std::ffi::c_void) -> i32; + + /// Launch a JIT trace-gen kernel matching the v1 signature: + /// (uint32_t* d_output, size_t H, int N, + /// const uint8_t* d_arena, uint32_t range_max_bits) + /// Synchronizes before returning. + pub fn powdr_nvrtc_launch_jit_v1( + function: *mut std::ffi::c_void, + d_output: *mut u32, + h: usize, + num_apc_calls: i32, + d_arena: *const u8, + range_max_bits: u32, + grid_x: u32, + block_x: u32, + ) -> i32; + + /// Upload `host_data` of `size_bytes` into a `__constant__` symbol + /// in the given module. Returns 0 on success. + pub fn powdr_nvrtc_set_constant_symbol( + module: *mut std::ffi::c_void, + sym_name: *const std::ffi::c_char, + host_data: *const std::ffi::c_void, + size_bytes: usize, + ) -> i32; + + /// Launch a per-APC codegen bus kernel (bus_v4) for var_range or tuple2. + /// No d_ops pointer (constants baked into source) and no n_ops arg. + /// `has_extra1=1` only for tuple2. + pub fn powdr_nvrtc_launch_bus_v4( + function: *mut std::ffi::c_void, + d_output: *const u32, + num_apc_calls: i32, + h: u64, + d_hist: *mut u32, + extra0: u32, + extra1: u32, + has_extra1: u32, + grid_x: u32, + block_x: u32, + ) -> i32; + + /// Launch a per-APC codegen bus kernel for bitwise (range or xor). + /// No d_ops, no n_ops, no extra args (size baked in). + pub fn powdr_nvrtc_launch_bus_v4_bitwise( + function: *mut std::ffi::c_void, + d_output: *const u32, + num_apc_calls: i32, + h: u64, + d_hist: *mut u32, + grid_x: u32, + block_x: u32, + ) -> i32; + + /// Launch a kind-templated bus kernel (bus_v3). Same shape as v2 plus a + /// `d_ops` pointer (op tables now in global memory). + pub fn powdr_nvrtc_launch_bus_v3( + function: *mut std::ffi::c_void, + d_output: *const u32, + num_apc_calls: i32, + h: u64, + d_hist: *mut u32, + extra0: u32, + extra1: u32, + has_extra1: u32, + n_ops: u32, + d_ops: *const std::ffi::c_void, + grid_x: u32, + block_x: u32, + ) -> i32; + + /// Launch a kind-templated bus kernel (bus_v2). `extra0`/`extra1` are + /// the kind-specific extra kernel args (e.g., `var_num_bins` for + /// var_range; tuple2 sizes for tuple2; nothing for bitwise — pass 0). + /// Set `has_extra1 = 1` only for tuple2 (which has two size args). + pub fn powdr_nvrtc_launch_bus_v2( + function: *mut std::ffi::c_void, + d_output: *const u32, + num_apc_calls: i32, + h: u64, + d_hist: *mut u32, + extra0: u32, + extra1: u32, + has_extra1: u32, + n_ops: u32, + grid_x: u32, + block_x: u32, + ) -> i32; + + /// Launch a JIT bus kernel matching the bus_v1 signature: + /// (const u32* d_output, int N, u64 H, + /// u32 var_range_bus_id, u32* d_var_hist, u32 var_num_bins, + /// u32 tuple2_bus_id, u32* d_tuple2_hist, u32 tuple2_sz0, u32 tuple2_sz1, + /// u32 bitwise_bus_id, u32* d_bitwise_hist) + /// Synchronizes before returning. + pub fn powdr_nvrtc_launch_bus_v1( + function: *mut std::ffi::c_void, + d_output: *const u32, + num_apc_calls: i32, + h: u64, + 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, + grid_x: u32, + block_x: u32, + ) -> i32; +} + +/// JIT column computation descriptor — matches CUDA JitColumnDesc struct. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct JitColumnDesc { + pub comp_type: u16, + pub apc_col: u16, + pub p: [u16; 6], +} + +/// JIT instruction descriptor — matches CUDA JitInstructionDesc struct. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct JitInstructionDesc { + pub arena_offset: u32, + pub record_stride: u32, + pub record_offset: u32, + pub col_desc_start: u32, + pub col_desc_count: u32, } #[repr(C)] @@ -221,3 +419,63 @@ pub fn apc_apply_bus( )) } } + +/// Phase 0 spike — same dispatch shape as `apc_apply_bus` but with bytecode +/// evaluation removed (one trace-cell load + one histogram update per row +/// per interaction). Output histograms are MEANINGLESS — this is a timing +/// instrument only. +#[allow(clippy::too_many_arguments)] +pub fn apc_apply_bus_spike( + output: &DeviceMatrix, + num_apc_calls: usize, + interactions: 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_spike( + output.buffer().as_ptr(), + num_apc_calls as i32, + interactions.as_ptr(), + interactions.len(), + var_range_bus_id, + var_range_count.as_mut_ptr() as *mut u32, + var_range_count.len(), + 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, + )) + } +} + +/// Safe wrapper for the JIT trace generation kernel. +pub fn apc_jit_tracegen( + output: &mut DeviceMatrix, + arena: &DeviceBuffer, + instructions: &DeviceBuffer, + col_descs: &DeviceBuffer, + num_apc_calls: usize, + range_max_bits: u32, +) -> Result<(), CudaError> { + let output_height = output.height(); + unsafe { + CudaError::from_result(_apc_jit_tracegen( + output.buffer().as_mut_ptr(), + output_height, + num_apc_calls as i32, + arena.as_ptr(), + instructions.as_ptr(), + instructions.len() as i32, + col_descs.as_ptr(), + range_max_bits, + )) + } +} diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs index 9a965f1fd7..44af69b947 100644 --- a/openvm/src/customize_exe.rs +++ b/openvm/src/customize_exe.rs @@ -16,7 +16,7 @@ use openvm_circuit::system::memory::online::GuestMemory; use openvm_instructions::instruction::Instruction as OpenVmInstruction; use openvm_instructions::program::DEFAULT_PC_STEP; use openvm_instructions::VmOpcode; -use openvm_stark_backend::p3_field::{FieldAlgebra, PrimeField32}; +use openvm_stark_backend::p3_field::{PrimeCharacteristicRing, PrimeField32}; use openvm_stark_sdk::p3_baby_bear::BabyBear; use powdr_autoprecompiles::adapter::{ Adapter, AdapterApc, AdapterApcWithStats, ApcWithStats, PgoAdapter, @@ -99,9 +99,7 @@ impl<'a, ISA: OpenVmISA> Adapter for BabyBearOpenVmApcAdapter<'a, ISA> { type ExecutionState = OpenVmExecutionState<'a, BabyBear, ISA>; fn into_field(e: Self::PowdrField) -> Self::Field { - openvm_stark_sdk::p3_baby_bear::BabyBear::from_canonical_u32( - e.to_integer().try_into_u32().unwrap(), - ) + openvm_stark_sdk::p3_baby_bear::BabyBear::from_u32(e.to_integer().try_into_u32().unwrap()) } fn from_field(e: Self::Field) -> Self::PowdrField { @@ -196,7 +194,7 @@ impl Instruction for Instr { ]; // The PC lookup row has the format: // [pc, opcode, a, b, c, d, e, f, g] - let pc = F::from_canonical_u32(pc.try_into().unwrap()); + let pc = F::from_u32(pc.try_into().unwrap()); once(pc).chain(args).collect() } } @@ -276,7 +274,7 @@ pub fn customize<'a, ISA: OpenVmISA, P: PgoAdapter( .enumerate() .collect::>(); - for (air_id, proving_context) in &ctx.per_air { - let main = proving_context.common_main.as_ref().unwrap(); + for (air_id, proving_context) in &ctx.per_trace { + let main = proving_context.common_main.clone(); let air_name = global_airs[air_id].name(); let Some(machine) = &airs.get_air_machine(&air_name) else { // air_name_to_machine only contains instruction AIRs, and we are only diff --git a/openvm/src/extraction_utils.rs b/openvm/src/extraction_utils.rs index 540a65cf5c..6137b34f09 100644 --- a/openvm/src/extraction_utils.rs +++ b/openvm/src/extraction_utils.rs @@ -7,18 +7,12 @@ use openvm_circuit::arch::{ AirInventory, AirInventoryError, ExecutorInventory, ExecutorInventoryError, SystemConfig, VmCircuitConfig, VmExecutionConfig, }; -use openvm_circuit::system::memory::interface::MemoryInterfaceAirs; use openvm_circuit_primitives::bitwise_op_lookup::SharedBitwiseOperationLookupChip; use openvm_circuit_primitives::range_tuple::SharedRangeTupleCheckerChip; use openvm_instructions::VmOpcode; use openvm_stark_backend::air_builders::symbolic::SymbolicRapBuilder; -use openvm_stark_backend::interaction::fri_log_up::find_interaction_chunks; use openvm_stark_backend::{ - air_builders::symbolic::SymbolicConstraints, config::StarkGenericConfig, rap::AnyRap, -}; -use openvm_stark_sdk::config::{ - baby_bear_poseidon2::{config_from_perm, default_perm}, - fri_params::SecurityParameters, + air_builders::symbolic::SymbolicConstraints, p3_air::BaseAir, AnyAir, StarkProtocolConfig, }; use openvm_stark_sdk::p3_baby_bear::{self, BabyBear}; use powdr_autoprecompiles::bus_map::BusType; @@ -42,9 +36,6 @@ use crate::utils::UnsupportedOpenVmReferenceError; use crate::AirMetrics; use crate::{air_builder::AirKeygenBuilder, BabyBearSC}; -// TODO: Use ` as FieldExtensionAlgebra>>::D` instead after fixing p3 dependency -const EXT_DEGREE: usize = 4; - #[derive(Clone, Serialize, Deserialize)] pub struct OriginalAirs { /// The degree bound used when building the airs @@ -344,15 +335,7 @@ impl OriginalVmConfig { BusType::ExecutionBridge, ), ( - // TODO: make getting memory bus index a helper function - match &memory_air.interface { - MemoryInterfaceAirs::Volatile { boundary } => { - boundary.memory_bus.inner.index - } - MemoryInterfaceAirs::Persistent { boundary, .. } => { - boundary.memory_bus.inner.index - } - }, + memory_air.interface.boundary.memory_bus.inner.index, BusType::Memory, ), (connector_air.program_bus.index(), BusType::PcLookup), @@ -395,8 +378,8 @@ impl OriginalVmConfig { } } -pub fn get_columns(air: Arc>) -> Vec> { - let width = air.width(); +pub fn get_columns(air: Arc>) -> Vec> { + let width = as BaseAir>::width(air.as_ref()); air.columns() .inspect(|columns| { assert_eq!(columns.len(), width); @@ -407,19 +390,19 @@ pub fn get_columns(air: Arc>) -> Vec> { .collect() } -pub fn get_name(air: Arc>) -> String { +pub fn get_name(air: Arc>) -> String { air.name() } pub fn get_constraints( - air: Arc>, + air: Arc>, ) -> SymbolicConstraints { let builder = symbolic_builder_with_degree(air, None); builder.constraints() } -pub fn get_air_metrics(air: Arc>, max_degree: usize) -> AirMetrics { - let main = air.width(); +pub fn get_air_metrics(air: Arc>, max_degree: usize) -> AirMetrics { + let main = as BaseAir>::width(air.as_ref()); let symbolic_rap_builder = symbolic_builder_with_degree(air, Some(max_degree)); let preprocessed = symbolic_rap_builder.width().preprocessed.unwrap_or(0); @@ -429,31 +412,18 @@ pub fn get_air_metrics(air: Arc>, max_degree: usize) -> A interactions, } = symbolic_rap_builder.constraints(); - let log_up = (find_interaction_chunks(&interactions, max_degree) - .interaction_partitions() - .len() - + 1) - * EXT_DEGREE; - AirMetrics { - widths: AirWidths { - preprocessed, - main, - log_up, - }, + widths: AirWidths { preprocessed, main }, constraints: constraints.len(), bus_interactions: interactions.len(), } } pub fn symbolic_builder_with_degree( - air: Arc>, + air: Arc>, max_constraint_degree: Option, ) -> SymbolicRapBuilder { - let perm = default_perm(); - let security_params = SecurityParameters::standard_fast(); - let config = config_from_perm(&perm, security_params); - let air_keygen_builder = AirKeygenBuilder::new(config.pcs(), air); + let air_keygen_builder = AirKeygenBuilder::new(air); air_keygen_builder.get_symbolic_builder(max_constraint_degree) } @@ -461,7 +431,6 @@ pub fn symbolic_builder_with_degree( pub struct AirWidths { pub preprocessed: usize, pub main: usize, - pub log_up: usize, } impl Add for AirWidths { @@ -470,7 +439,6 @@ impl Add for AirWidths { AirWidths { preprocessed: self.preprocessed + rhs.preprocessed, main: self.main + rhs.main, - log_up: self.log_up + rhs.log_up, } } } @@ -481,7 +449,6 @@ impl Sub for AirWidths { AirWidths { preprocessed: self.preprocessed - rhs.preprocessed, main: self.main - rhs.main, - log_up: self.log_up - rhs.log_up, } } } @@ -494,7 +461,7 @@ impl Sum for AirWidths { impl AirWidths { pub fn total(&self) -> usize { - self.preprocessed + self.main + self.log_up + self.preprocessed + self.main } } @@ -502,11 +469,10 @@ impl std::fmt::Display for AirWidths { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "Total Width: {} (Preprocessed: {} Main: {}, Log Up: {})", - self.preprocessed + self.main + self.log_up, + "Total Width: {} (Preprocessed: {}, Main: {})", + self.preprocessed + self.main, self.preprocessed, self.main, - self.log_up ) } } diff --git a/openvm/src/isa.rs b/openvm/src/isa.rs index d18cb40334..ed9838c83e 100644 --- a/openvm/src/isa.rs +++ b/openvm/src/isa.rs @@ -4,19 +4,20 @@ use std::sync::Arc; use openvm_circuit::arch::{ AirInventory, AirInventoryError, AnyEnum, ChipInventory, ChipInventoryError, DenseRecordArena, Executor, InterpreterExecutor, MatrixRecordArena, MeteredExecutor, PreflightExecutor, - VmBuilder, VmChipComplex, VmCircuitExtension, VmConfig, VmExecutionConfig, + VmBuilder, VmChipComplex, VmCircuitExtension, VmConfig, VmExecutionConfig, VmField, }; #[cfg(feature = "cuda")] use openvm_circuit::system::cuda::SystemChipInventoryGPU; use openvm_circuit::system::SystemChipInventory; +use openvm_cpu_backend::CpuBackend; #[cfg(feature = "cuda")] -use openvm_cuda_backend::engine::GpuBabyBearPoseidon2Engine; +use openvm_cuda_backend::BabyBearPoseidon2GpuEngine as GpuBabyBearPoseidon2CpuEngine; #[cfg(feature = "cuda")] -use openvm_cuda_backend::prover_backend::GpuBackend; +use openvm_cuda_backend::GpuBackend; use openvm_instructions::{instruction::Instruction, VmOpcode}; -use openvm_sdk::config::TranspilerConfig; -use openvm_stark_backend::{config::Val, p3_field::PrimeField32, prover::cpu::CpuBackend}; -use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine; +use openvm_sdk_config::TranspilerConfig; +use openvm_stark_backend::{p3_field::PrimeField32, Val}; +use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2CpuEngine; use openvm_stark_sdk::p3_baby_bear::BabyBear; use powdr_riscv_elf::debug_info::SymbolTable; @@ -47,7 +48,7 @@ pub trait OpenVmISA: Send + Sync + Clone + 'static + Default { /// The original linked program, for example, an elf for riscv. It must allow recovering the jump destinations. type LinkedProgram<'a>; - type Executor: AnyEnum + type Executor: AnyEnum + InterpreterExecutor + Executor + MeteredExecutor @@ -64,7 +65,7 @@ pub trait OpenVmISA: Send + Sync + Clone + 'static + Default { type CpuBuilder: Clone + Default + VmBuilder< - BabyBearPoseidon2Engine, + BabyBearPoseidon2CpuEngine, VmConfig = Self::Config, SystemChipInventory = SystemChipInventory, RecordArena = MatrixRecordArena>, @@ -74,7 +75,7 @@ pub trait OpenVmISA: Send + Sync + Clone + 'static + Default { type GpuBuilder: Clone + Default + VmBuilder< - GpuBabyBearPoseidon2Engine, + GpuBabyBearPoseidon2CpuEngine, VmConfig = Self::Config, SystemChipInventory = SystemChipInventoryGPU, RecordArena = DenseRecordArena, diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs index 7d76028643..f0f57aaccb 100644 --- a/openvm/src/lib.rs +++ b/openvm/src/lib.rs @@ -7,7 +7,7 @@ use openvm_circuit::arch::{ AirInventory, AirInventoryError, ChipInventory, ChipInventoryError, ExecutorInventory, ExecutorInventoryError, InitFileGenerator, MatrixRecordArena, RowMajorMatrixArena, SystemConfig, VmBuilder, VmChipComplex, VmCircuitConfig, VmCircuitExtension, VmExecutionConfig, - VmProverExtension, + VmField, VmProverExtension, }; use openvm_circuit::system::SystemChipInventory; use openvm_circuit::{circuit_derive::Chip, derive::AnyEnum}; @@ -15,20 +15,18 @@ use openvm_circuit_derive::{ AotExecutor, AotMeteredExecutor, Executor, MeteredExecutor, PreflightExecutor, }; -use openvm_sdk::config::TranspilerConfig; -use openvm_sdk::GenericSdk; +use openvm_cpu_backend::{CpuBackend, CpuDevice}; use openvm_sdk::{ - config::{AppConfig, DEFAULT_APP_LOG_BLOWUP}, - StdIn, + config::{AggregationSystemParams, AppConfig, DEFAULT_APP_LOG_BLOWUP}, + GenericSdk, StdIn, }; -use openvm_stark_backend::config::{StarkGenericConfig, Val}; -use openvm_stark_backend::engine::StarkEngine; -use openvm_stark_backend::prover::cpu::{CpuBackend, CpuDevice}; -use openvm_stark_backend::prover::hal::ProverBackend; -use openvm_stark_sdk::config::{ - baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine}, - FriParameters, +use openvm_sdk_config::TranspilerConfig; +use openvm_stark_backend::prover::ProverBackend; +use openvm_stark_backend::{StarkEngine, StarkProtocolConfig, Val}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{ + BabyBearPoseidon2Config, BabyBearPoseidon2CpuEngine, }; +use openvm_stark_sdk::config::{app_params_with_100_bits_security, MAX_APP_LOG_STACKED_HEIGHT}; use openvm_stark_sdk::openvm_stark_backend::p3_field::PrimeField32; use openvm_stark_sdk::p3_baby_bear::BabyBear; use openvm_transpiler::transpiler::Transpiler; @@ -70,14 +68,12 @@ pub type BabyBearSC = BabyBearPoseidon2Config; cfg_if::cfg_if! { if #[cfg(feature = "cuda")] { - pub use openvm_cuda_backend::engine::GpuBabyBearPoseidon2Engine; - pub use openvm_native_circuit::NativeGpuBuilder; - pub type PowdrSdkGpu = GenericSdk, NativeGpuBuilder>; - pub type PowdrExecutionProfileSdkGpu = GenericSdk::GpuBuilder, NativeGpuBuilder>; + pub use openvm_cuda_backend::BabyBearPoseidon2GpuEngine as GpuBabyBearPoseidon2CpuEngine; + pub type PowdrSdkGpu = GenericSdk>; + pub type PowdrExecutionProfileSdkGpu = GenericSdk::GpuBuilder>; pub use openvm_circuit::system::cuda::{extensions::SystemGpuBuilder, SystemChipInventoryGPU}; - pub use openvm_sdk::config::SdkVmGpuBuilder; - pub use openvm_cuda_backend::prover_backend::GpuBackend; + pub use openvm_cuda_backend::GpuBackend; pub use openvm_circuit_primitives::bitwise_op_lookup::BitwiseOperationLookupChipGPU; pub use openvm_circuit_primitives::range_tuple::RangeTupleCheckerChipGPU; pub use openvm_circuit_primitives::var_range::VariableRangeCheckerChipGPU; @@ -93,11 +89,10 @@ use openvm_circuit_primitives::range_tuple::{RangeTupleCheckerAir, SharedRangeTu use openvm_circuit_primitives::var_range::{ SharedVariableRangeCheckerChip, VariableRangeCheckerAir, }; -use openvm_native_circuit::NativeCpuBuilder; pub type PowdrSdkCpu = - GenericSdk, NativeCpuBuilder>; + GenericSdk>; pub type PowdrExecutionProfileSdkCpu = - GenericSdk::CpuBuilder, NativeCpuBuilder>; + GenericSdk::CpuBuilder>; pub const DEFAULT_OPENVM_DEGREE_BOUND: usize = 2 * DEFAULT_APP_LOG_BLOWUP + 1; pub const DEFAULT_DEGREE_BOUND: DegreeBound = DegreeBound { @@ -142,7 +137,7 @@ pub struct SpecializedConfigGpuBuilder { } #[cfg(feature = "cuda")] -impl VmBuilder for SpecializedConfigGpuBuilder { +impl VmBuilder for SpecializedConfigGpuBuilder { type VmConfig = SpecializedConfig; type SystemChipInventory = SystemChipInventoryGPU; type RecordArena = DenseRecordArena; @@ -155,13 +150,13 @@ impl VmBuilder for SpecializedConfig VmChipComplex, ChipInventoryError, > { - let mut chip_complex = VmBuilder::::create_chip_complex( + let mut chip_complex = VmBuilder::::create_chip_complex( &::GpuBuilder::default(), &config.original.config, circuit, )?; let inventory = &mut chip_complex.inventory; - VmProverExtension::::extend_prover( + VmProverExtension::::extend_prover( &PowdrGpuProverExt::::default(), &config.powdr, inventory, @@ -220,8 +215,11 @@ struct PowdrGpuProverExt { #[cfg(feature = "cuda")] impl - VmProverExtension> - for PowdrGpuProverExt + VmProverExtension< + GpuBabyBearPoseidon2CpuEngine, + DenseRecordArena, + PowdrExtension, + > for PowdrGpuProverExt { fn extend_prover( &self, @@ -338,7 +336,7 @@ where // This is the most robust method because bus ids are assigned at air creation time. fn get_periphery_bus_ids(inventory: &ChipInventory) -> PeripheryBusIds where - SC: StarkGenericConfig, + SC: StarkProtocolConfig, PB: ProverBackend, { let air_inventory = inventory.airs(); @@ -404,7 +402,7 @@ impl AsMut for SpecializedConfig { #[derive( AnyEnum, Chip, Executor, MeteredExecutor, AotExecutor, AotMeteredExecutor, PreflightExecutor, )] -pub enum SpecializedExecutor { +pub enum SpecializedExecutor { #[any_enum] OriginalExecutor(ISA::Executor), #[any_enum] @@ -412,9 +410,7 @@ pub enum SpecializedExecutor { } // We implement `From` by hand because we cannot prove that `ISA::Executor != PowdrExtensionExecutor` -impl From> - for SpecializedExecutor -{ +impl From> for SpecializedExecutor { fn from(value: PowdrExtensionExecutor) -> Self { Self::PowdrExecutor(value) } @@ -509,7 +505,7 @@ impl CompiledProgram { let air_inventory = self.vm_config.create_airs().unwrap(); let chip_complex = as VmBuilder< - BabyBearPoseidon2Engine, + BabyBearPoseidon2CpuEngine, >>::create_chip_complex( &SpecializedConfigCpuBuilder::default(), &self.vm_config, @@ -531,9 +527,6 @@ impl CompiledProgram { (Vec::new(), Vec::new()), |(mut powdr_air_metrics, mut non_powdr_air_metrics), air| { let name = air.name(); - // We actually give name "powdr_air_for_opcode_" to the AIRs, - // but OpenVM uses the actual Rust type (PowdrAir) as the name in this method. - // TODO this is hacky but not sure how to do it better rn. if name.starts_with("PowdrAir") { powdr_air_metrics.push(( get_air_metrics(air.clone(), max_degree), @@ -556,15 +549,14 @@ pub fn execute( let CompiledProgram { exe, vm_config } = program; // Set app configuration - let app_fri_params = - FriParameters::standard_with_100_bits_conjectured_security(DEFAULT_APP_LOG_BLOWUP); - let app_config = AppConfig::new(app_fri_params, vm_config.clone()); + let system_params = app_params_with_100_bits_security(MAX_APP_LOG_STACKED_HEIGHT); + let app_config = AppConfig::new(vm_config.clone(), system_params); // prepare for execute #[cfg(feature = "cuda")] - let sdk = PowdrSdkGpu::new(app_config).unwrap(); + let sdk = PowdrSdkGpu::new(app_config, AggregationSystemParams::default()).unwrap(); #[cfg(not(feature = "cuda"))] - let sdk = PowdrSdkCpu::new(app_config).unwrap(); + let sdk = PowdrSdkCpu::new(app_config, AggregationSystemParams::default()).unwrap(); let output = sdk.execute(exe.clone(), inputs.clone()).unwrap(); @@ -582,12 +574,13 @@ pub fn execution_profile_from_guest( let program = Prog::from(&exe.program); // Set app configuration - let app_fri_params = - FriParameters::standard_with_100_bits_conjectured_security(DEFAULT_APP_LOG_BLOWUP); - let app_config = AppConfig::new(app_fri_params, vm_config.clone().config); + let system_params = app_params_with_100_bits_security(MAX_APP_LOG_STACKED_HEIGHT); + let app_config = AppConfig::new(vm_config.clone().config, system_params); // prepare for execute - let sdk = PowdrExecutionProfileSdkCpu::::new(app_config).unwrap(); + let sdk = + PowdrExecutionProfileSdkCpu::::new(app_config, AggregationSystemParams::default()) + .unwrap(); execution_profile::>(&program, || { sdk.execute_interpreted(exe.clone(), inputs.clone()) diff --git a/openvm/src/powdr_extension/chip.rs b/openvm/src/powdr_extension/chip.rs index 4227c58276..eaef88fb7c 100644 --- a/openvm/src/powdr_extension/chip.rs +++ b/openvm/src/powdr_extension/chip.rs @@ -16,14 +16,12 @@ use itertools::Itertools; use openvm_circuit::arch::MatrixRecordArena; use openvm_stark_backend::{ p3_air::{Air, BaseAir}, - rap::ColumnsAir, + ColumnsAir, }; use openvm_stark_backend::{ - interaction::InteractionBuilder, - p3_field::PrimeField32, - p3_matrix::Matrix, - rap::{BaseAirWithPublicValues, PartitionedBaseAir}, + interaction::InteractionBuilder, p3_field::PrimeField32, p3_matrix::Matrix, + BaseAirWithPublicValues, PartitionedBaseAir, }; use openvm_stark_sdk::p3_baby_bear::BabyBear; use powdr_autoprecompiles::{ @@ -99,13 +97,14 @@ where { fn eval(&self, builder: &mut AB) { let main = builder.main(); - let witnesses = main.row_slice(0); + let witnesses = main.row_slice(0).expect("row_slice(0) should exist"); // TODO: cache? + let witness_slice: &[AB::Var] = &witnesses; let witness_values: BTreeMap = self .columns .iter() .map(|c| c.id) - .zip_eq(witnesses.iter().cloned()) + .zip_eq(witness_slice.iter().copied()) .collect(); let witness_evaluator = WitnessEvaluator::new(&witness_values); diff --git a/openvm/src/powdr_extension/executor/mod.rs b/openvm/src/powdr_extension/executor/mod.rs index e8487beacb..7aaeaaf01c 100644 --- a/openvm/src/powdr_extension/executor/mod.rs +++ b/openvm/src/powdr_extension/executor/mod.rs @@ -166,6 +166,13 @@ impl InitializedOriginalArenas { let index = *self.air_name_to_arena_index.get(air_name)?; self.arenas[index].take().map(|arena_pair| arena_pair.real) } + + /// Borrow a real arena without consuming it. Returns None if the air_name + /// is not found or the arena has already been taken. + pub fn peek_real_arena(&self, air_name: &str) -> Option<&A> { + let index = *self.air_name_to_arena_index.get(air_name)?; + self.arenas[index].as_ref().map(|pair| &pair.real) + } } pub struct ArenaPair { @@ -468,7 +475,6 @@ impl PreflightExecutor> memory, streams, rng, - custom_pvs, ctx, #[cfg(feature = "metrics")] metrics, @@ -507,7 +513,6 @@ impl PreflightExecutor> memory, streams, rng, - custom_pvs, // We execute in the context of the relevant original table ctx: ctx_arena, // TODO: should we pass around the same metrics object, or snapshot it at the beginning of this method and apply a single update at the end? @@ -542,7 +547,6 @@ impl PreflightExecutor for PowdrExec memory, streams, rng, - custom_pvs, ctx, #[cfg(feature = "metrics")] metrics, @@ -581,7 +585,6 @@ impl PreflightExecutor for PowdrExec memory, streams, rng, - custom_pvs, // We execute in the context of the relevant original table ctx: ctx_arena, // TODO: should we pass around the same metrics object, or snapshot it at the beginning of this method and apply a single update at the end? diff --git a/openvm/src/powdr_extension/trace_generator/common.rs b/openvm/src/powdr_extension/trace_generator/common.rs index 9cf0e0c641..b15a941d91 100644 --- a/openvm/src/powdr_extension/trace_generator/common.rs +++ b/openvm/src/powdr_extension/trace_generator/common.rs @@ -1,14 +1,14 @@ use derive_more::From; +use openvm_circuit::arch::VmField; use openvm_circuit::system::phantom::PhantomExecutor; use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor}; use openvm_circuit_primitives::Chip; -use openvm_stark_backend::p3_field::PrimeField32; use crate::isa::OpenVmISA; #[allow(clippy::large_enum_variant)] #[derive(Chip, PreflightExecutor, Executor, MeteredExecutor, AnyEnum)] -pub enum DummyExecutor { +pub enum DummyExecutor { #[any_enum] Base(ISA::Executor), #[any_enum] @@ -16,6 +16,6 @@ pub enum DummyExecutor { } #[derive(Chip, PreflightExecutor, Executor, MeteredExecutor, From, AnyEnum)] -pub enum SharedExecutor { +pub enum SharedExecutor { Phantom(PhantomExecutor), } diff --git a/openvm/src/powdr_extension/trace_generator/cpu/inventory.rs b/openvm/src/powdr_extension/trace_generator/cpu/inventory.rs index ef77e4abd3..e9674fc6a5 100644 --- a/openvm/src/powdr_extension/trace_generator/cpu/inventory.rs +++ b/openvm/src/powdr_extension/trace_generator/cpu/inventory.rs @@ -2,7 +2,8 @@ use openvm_circuit::{ arch::{MatrixRecordArena, VmChipComplex}, system::SystemChipInventory, }; -use openvm_stark_backend::{config::Val, prover::cpu::CpuBackend}; +use openvm_cpu_backend::CpuBackend; +use openvm_stark_backend::Val; /// A dummy inventory used for execution of autoprecompiles /// It extends the `SdkVmConfigExecutor` and `SdkVmConfigPeriphery`, providing them with shared, pre-loaded periphery chips to avoid memory allocations by each SDK chip diff --git a/openvm/src/powdr_extension/trace_generator/cpu/mod.rs b/openvm/src/powdr_extension/trace_generator/cpu/mod.rs index 78da83a182..0351b5eac5 100644 --- a/openvm/src/powdr_extension/trace_generator/cpu/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/cpu/mod.rs @@ -2,16 +2,19 @@ use std::{collections::HashMap, sync::Arc}; use itertools::Itertools; use openvm_circuit::{arch::MatrixRecordArena, utils::next_power_of_two_or_zero}; +use openvm_circuit_primitives::Chip; +use openvm_stark_backend::p3_maybe_rayon::prelude::*; use openvm_stark_backend::{ - p3_field::{Field, FieldAlgebra, PrimeField32}, + p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, p3_matrix::dense::{DenseMatrix, RowMajorMatrix}, - prover::{hal::ProverBackend, types::AirProvingContext}, - Chip, + prover::{AirProvingContext, ProverBackend}, }; use openvm_stark_sdk::p3_baby_bear::BabyBear; use powdr_autoprecompiles::trace_handler::TraceTrait; +use powdr_autoprecompiles::InstructionHandler; use powdr_constraint_solver::constraint_system::ComputationMethod; +use super::jit_mapping::{self, AirColumnMapping, ColumnComputation}; use crate::{ extraction_utils::{OriginalAirs, OriginalVmConfig}, isa::IsaApc, @@ -52,17 +55,44 @@ impl From>> for SharedCpuTrace { } } -impl>>, ISA: OpenVmISA> Chip +impl>, ISA: OpenVmISA> Chip for PowdrChipCpu { fn generate_proving_ctx(&self, _: R) -> AirProvingContext { tracing::trace!("Generating air proof input for PowdrChip {}", self.name); - let trace = self - .trace_generator - .generate_witness(self.record_arena_by_air_name.take()); + let use_jit = std::env::var("POWDR_JIT_TRACEGEN").is_ok(); + let validate_jit = std::env::var("POWDR_JIT_VALIDATE").is_ok(); - AirProvingContext::simple(Arc::new(trace), vec![]) + let row_major = if use_jit { + let arenas = self.record_arena_by_air_name.take(); + match self.trace_generator.generate_witness_jit(arenas) { + Ok(trace) => { + tracing::info!("JIT trace gen used for PowdrChip {}", self.name); + trace + } + Err(arenas) => { + tracing::warn!( + "JIT trace gen not available for PowdrChip {}, falling back to standard path", + self.name + ); + self.trace_generator.generate_witness(arenas) + } + } + } else if validate_jit { + // Validation mode: run standard path first, then run JIT path on a + // clone of the arenas and compare the APC columns that JIT fills. + // This validates the JIT mapping is correct without requiring all + // AIR types to be supported. + self.trace_generator.generate_witness_and_validate_jit( + self.record_arena_by_air_name.take(), + ) + } else { + self.trace_generator + .generate_witness(self.record_arena_by_air_name.take()) + }; + + AirProvingContext::simple(row_major, vec![]) } } @@ -119,6 +149,8 @@ impl PowdrTraceGeneratorCpu { .inventory }; + let jit_debug = std::env::var("POWDR_JIT_DEBUG").is_ok(); + let dummy_trace_by_air_name: HashMap> = chip_inventory .chips() .iter() @@ -134,9 +166,62 @@ impl PowdrTraceGeneratorCpu { } }; - let shared_trace = chip.generate_proving_ctx(record_arena).common_main.unwrap(); + if jit_debug { + // Dump arena layout before fill_trace_row consumes it + let arena_width = record_arena.width; + let arena_offset = record_arena.trace_offset; + let rows_used = arena_offset / arena_width; + + tracing::info!( + "JIT_DEBUG: AIR '{}' arena: width={}, rows_used={}, total_values={}", + air_name, arena_width, rows_used, record_arena.trace_buffer.len() + ); + + // Dump first row's raw record bytes + if rows_used > 0 { + let row0 = &record_arena.trace_buffer[..arena_width]; + // Print raw u32 values (NOT Montgomery-decoded) — these are record struct bytes + let row0_raw: Vec = row0 + .iter() + .map(|v| { + let bytes: [u8; 4] = unsafe { std::mem::transmute_copy(v) }; + u32::from_le_bytes(bytes) + }) + .collect(); + tracing::info!( + "JIT_DEBUG: row0 raw bytes (as u32 LE): {:?}", + row0_raw + ); + // Also print as hex for easier struct field reading + let row0_hex: Vec = row0_raw.iter().map(|v| format!("{:#010x}", v)).collect(); + tracing::info!( + "JIT_DEBUG: row0 raw bytes (hex): {:?}", + row0_hex + ); + } + } - Some((air_name, SharedCpuTrace::from(shared_trace))) + let row_major_trace = chip.generate_proving_ctx(record_arena).common_main; + + if jit_debug { + // Dump first row after fill_trace_row + let trace_width = row_major_trace.width; + let trace_height = row_major_trace.values.len() / trace_width; + tracing::info!( + "JIT_DEBUG: AIR '{}' trace: width={}, height={}", + air_name, trace_width, trace_height + ); + if trace_height > 0 { + let row0 = &row_major_trace.values[..trace_width]; + let row0_u32: Vec = row0.iter().map(|v| v.as_canonical_u32()).collect(); + tracing::info!( + "JIT_DEBUG: row0 after fill_trace (as u32): {:?}", + row0_u32 + ); + } + } + + Some((air_name, SharedCpuTrace::from(Arc::new(row_major_trace)))) }) .collect(); @@ -152,21 +237,42 @@ impl PowdrTraceGeneratorCpu { &self.apc, ); - // allocate for apc trace + // Build dense Vec indexed by poly ID for O(1) column lookups in the hot loop. + // Poly IDs may be sparse (gaps between IDs), so the Vec is sized to max_id + 1. let width = apc_poly_id_to_index.len(); + let max_poly_id = apc_poly_id_to_index.keys().last().copied().unwrap_or(0) as usize; + let apc_poly_id_to_index: Vec = (0..=max_poly_id) + .map(|id| apc_poly_id_to_index.get(&(id as u64)).copied().unwrap_or(0)) + .collect(); + + // Compile bus interactions once before the hot loop + let compiled_interactions = { + use powdr_autoprecompiles::expression::CompiledBusInteraction; + CompiledBusInteraction::compile_all( + &self.apc.machine().bus_interactions, + &apc_poly_id_to_index, + BabyBear::ZERO, + BabyBear::ONE, + ) + }; + + // allocate for apc trace let height = next_power_of_two_or_zero(num_apc_calls); - let mut values = ::zero_vec(height * width); + let mut values = ::zero_vec(height * width); + + // Go through the final table and fill in the values. + // Parallelized: the original code used chunks_mut (serial) because the + // periphery.apply() calls were not thread-safe. Now that periphery chips + // use AtomicU32 counters, we can use par_chunks_mut safely. + // Extract references to avoid capturing `self` (ISA::Config is not Sync). + let periphery_real = &self.periphery.real; + let periphery_bus_ids = &self.periphery.bus_ids; - // go through the final table and fill in the values values - // a record is `width` values - // TODO: optimize by parallelizing on chunks of rows, currently fails because `dyn AnyChip>>` is not `Send` - .chunks_mut(width) - .zip(dummy_values) + .par_chunks_mut(width) + .zip(dummy_values.into_par_iter()) .for_each(|(row_slice, dummy_values)| { - // map the dummy rows to the autoprecompile row - - use powdr_autoprecompiles::expression::MappingRowEvaluator; + // Copy dummy rows to APC row for (dummy_row, dummy_trace_index_to_apc_index) in dummy_values .iter() .map(|r| &r.data[r.start..r.start + r.length]) @@ -177,53 +283,516 @@ impl PowdrTraceGeneratorCpu { } } - // Fill in the columns we have to compute from other columns - // (these are either new columns or for example the "is_valid" column). + // Compute derived columns for derived_column in columns_to_compute { - let col_index = apc_poly_id_to_index[&derived_column.variable.id]; + let col_index = apc_poly_id_to_index[derived_column.variable.id as usize]; row_slice[col_index] = match &derived_column.computation_method { ComputationMethod::Constant(c) => *c, ComputationMethod::QuotientOrZero(e1, e2) => { use powdr_number::ExpressionConvertible; let divisor_val = e2.to_expression(&|n| *n, &|column_ref| { - row_slice[apc_poly_id_to_index[&column_ref.id]] + row_slice[apc_poly_id_to_index[column_ref.id as usize]] }); if divisor_val.is_zero() { BabyBear::ZERO } else { divisor_val.inverse() * e1.to_expression(&|n| *n, &|column_ref| { - row_slice[apc_poly_id_to_index[&column_ref.id]] + row_slice[apc_poly_id_to_index[column_ref.id as usize]] }) } } }; } - let evaluator = MappingRowEvaluator::new(row_slice, &apc_poly_id_to_index); + // Evaluate bus interactions using compiled expressions. + // Periphery chips use AtomicU32 counters — thread-safe. + for ci in &compiled_interactions { + let mult = ci.mult.eval(row_slice); + periphery_real.apply( + ci.id as u16, + mult.as_canonical_u32(), + ci.args.iter().map(|a| a.eval(row_slice).as_canonical_u32()), + periphery_bus_ids, + ); + } + }); - // replay the side effects of this row on the main periphery - self.apc - .machine() - .bus_interactions - .iter() - .for_each(|interaction| { - use powdr_autoprecompiles::expression::{ - AlgebraicEvaluator, ConcreteBusInteraction, - }; - - let ConcreteBusInteraction { id, mult, args } = - evaluator.eval_bus_interaction(interaction); - self.periphery.real.apply( - id as u16, - mult.as_canonical_u32(), - args.map(|arg| arg.as_canonical_u32()), - &self.periphery.bus_ids, + RowMajorMatrix::new(values, width) + } + + /// JIT trace generation: reads record bytes directly from arenas, + /// computes only surviving APC columns, bypassing fill_trace_row entirely. + /// + /// Returns `Ok(trace)` if JIT was used, `Err(arenas)` if unsupported + /// (arenas are returned unconsumed for fallback to standard path). + pub fn generate_witness_jit( + &self, + original_arenas: OriginalArenas>, + ) -> Result, OriginalArenas>> { + use std::collections::BTreeMap; + + let width = self.apc.machine().main_columns().count(); + + let mut original_arenas = match original_arenas { + OriginalArenas::Initialized(arenas) => arenas, + OriginalArenas::Uninitialized => { + return Ok(RowMajorMatrix::new(vec![], width)); + } + }; + + let num_apc_calls = original_arenas.number_of_calls; + let jit_debug = std::env::var("POWDR_JIT_DEBUG").is_ok(); + + // Build the same instruction metadata as generate_trace, but we'll use + // column mappings instead of dummy traces. + + // Keep only instructions with surviving substitutions + let instructions_with_subs: Vec<_> = self + .apc + .instructions() + .zip_eq(self.apc.subs().iter()) + .filter(|(_, subs)| !subs.is_empty()) + .collect(); + + // Get AIR names for each instruction + let air_names: Vec = instructions_with_subs + .iter() + .map(|(instr, _)| { + let (air_id, _) = self.original_airs.get_instruction_air_and_id(instr); + air_id + }) + .collect(); + + // Count occurrences per AIR (= instructions of this type per APC call) + let air_id_occurrences: HashMap<&str, usize> = air_names + .iter() + .map(|s| s.as_str()) + .counts(); + + // Compute per-instruction offset within its AIR type + let instruction_offsets: Vec = air_names + .iter() + .scan( + HashMap::<&str, usize>::default(), + |counts, air_name| { + let count = counts.entry(air_name.as_str()).or_default(); + let current = *count; + *count += 1; + Some(current) + }, + ) + .collect(); + + // Build APC poly_id to index mapping + let apc_poly_id_to_index: BTreeMap = self + .apc + .machine() + .main_columns() + .enumerate() + .map(|(index, c)| (c.id, index)) + .collect(); + + // Build substitution mapping: for each instruction, (original_poly_index -> apc_col_index) + let subs_by_instruction: Vec> = instructions_with_subs + .iter() + .map(|(_, subs)| { + subs.iter() + .map(|s| (s.original_poly_index, apc_poly_id_to_index[&s.apc_poly_id])) + .collect() + }) + .collect(); + + // Look up the column mapping for each AIR type + let mappings_by_air: HashMap<&str, AirColumnMapping> = { + let mut m = HashMap::new(); + m.insert( + "VmAirWrapper", + jit_mapping::base_alu_mapping(), + ); + m.insert( + "VmAirWrapper", + jit_mapping::loadstore_mapping(), + ); + m.insert( + "VmAirWrapper", + jit_mapping::shift_mapping(), + ); + m.insert( + "VmAirWrapper", + jit_mapping::branch_equal_mapping(), + ); + m + }; + + // Pre-check: verify ALL AIR types have mappings BEFORE consuming any arenas + for air_name in air_names.iter().unique() { + if !mappings_by_air.contains_key(air_name.as_str()) { + if jit_debug { + tracing::warn!( + "JIT_DEBUG: No JIT mapping for AIR '{}', falling back to standard path", + air_name + ); + } + // Wrap arenas back into OriginalArenas::Initialized and return + return Err(OriginalArenas::Initialized(original_arenas)); + } + } + + // All AIR types are supported. Now take arenas. + let mut arena_bytes_by_air: HashMap, usize)> = HashMap::new(); + + for air_name in air_names.iter().unique() { + if let Some(arena) = original_arenas.take_real_arena(air_name) { + let arena_width = arena.width; + // Convert trace_buffer Vec to Vec by reinterpreting + let byte_vec: Vec = unsafe { + let ptr = arena.trace_buffer.as_ptr() as *const u8; + let len = arena.trace_buffer.len() * 4; + std::slice::from_raw_parts(ptr, len).to_vec() + }; + arena_bytes_by_air.insert(air_name.clone(), (byte_vec, arena_width)); + } + } + + // Build per-instruction JIT info + // (We can't use InstructionJitInfo with lifetimes easily, so flatten the data) + struct FlatJitInfo { + arena_key: String, + occurrences: usize, + offset: usize, + #[allow(dead_code)] + mapping_air_name: &'static str, + } + + let jit_infos: Vec = air_names + .iter() + .zip(instruction_offsets.iter()) + .map(|(air_name, offset)| { + let occurrences = *air_id_occurrences.get(air_name.as_str()).unwrap(); + let mapping = mappings_by_air.get(air_name.as_str()).unwrap(); + FlatJitInfo { + arena_key: air_name.clone(), + occurrences, + offset: *offset, + mapping_air_name: mapping.air_name, + } + }) + .collect(); + + // Get range_max_bits from the periphery's range checker + let range_max_bits = self.periphery.real.range_checker.air.bus.range_max_bits as u32; + + if jit_debug { + tracing::info!( + "JIT_DEBUG: Starting JIT trace gen: {} APC calls, {} instructions with subs, {} APC cols, range_max_bits={}", + num_apc_calls, instructions_with_subs.len(), width, range_max_bits + ); + } + + // Build dense apc_poly_id_to_index for hot loop + let max_poly_id = apc_poly_id_to_index.keys().last().copied().unwrap_or(0) as usize; + let apc_poly_id_to_index_dense: Vec = (0..=max_poly_id) + .map(|id| apc_poly_id_to_index.get(&(id as u64)).copied().unwrap_or(0)) + .collect(); + + // Compile bus interactions + let compiled_interactions = { + use powdr_autoprecompiles::expression::CompiledBusInteraction; + CompiledBusInteraction::compile_all( + &self.apc.machine().bus_interactions, + &apc_poly_id_to_index_dense, + BabyBear::ZERO, + BabyBear::ONE, + ) + }; + + let columns_to_compute = &self.apc.machine().derived_columns; + + // Allocate APC trace + let height = next_power_of_two_or_zero(num_apc_calls); + let mut values = ::zero_vec(height * width); + + let periphery_real = &self.periphery.real; + let periphery_bus_ids = &self.periphery.bus_ids; + + // Process each APC row + values + .par_chunks_mut(width) + .enumerate() + .for_each(|(row_idx, row_slice)| { + if row_idx >= num_apc_calls { + return; // padding rows stay zero + } + + // For each instruction with surviving subs + for (_instr_idx, (jit_info, subs)) in + jit_infos.iter().zip(subs_by_instruction.iter()).enumerate() + { + let (arena_bytes, arena_width) = { + let (bytes, w) = arena_bytes_by_air.get(&jit_info.arena_key).unwrap(); + (bytes.as_slice(), *w) + }; + let mapping = mappings_by_air.get(jit_info.arena_key.as_str()).unwrap(); + + // Compute the record byte offset in the arena + let arena_row_idx = + row_idx * jit_info.occurrences + jit_info.offset; + let row_bytes = arena_width * 4; + let record_byte_start = arena_row_idx * row_bytes; + let record_byte_end = record_byte_start + row_bytes; + let record_bytes = &arena_bytes[record_byte_start..record_byte_end]; + + // For each surviving substitution, compute the column value + for &(original_poly_index, apc_col_index) in subs { + // Look up the computation for this column + let col_mapping = &mapping.columns[original_poly_index]; + debug_assert_eq!(col_mapping.col_index, original_poly_index); + + let value: BabyBear = jit_mapping::eval_column( + &col_mapping.computation, + record_bytes, + range_max_bits, ); - }); + row_slice[apc_col_index] = value; + } + } + + // Compute derived columns (same as standard path) + for derived_column in columns_to_compute { + let col_index = + apc_poly_id_to_index_dense[derived_column.variable.id as usize]; + row_slice[col_index] = match &derived_column.computation_method { + ComputationMethod::Constant(c) => *c, + ComputationMethod::QuotientOrZero(e1, e2) => { + use powdr_number::ExpressionConvertible; + let divisor_val = e2.to_expression(&|n| *n, &|column_ref| { + row_slice + [apc_poly_id_to_index_dense[column_ref.id as usize]] + }); + if divisor_val.is_zero() { + BabyBear::ZERO + } else { + divisor_val.inverse() + * e1.to_expression(&|n| *n, &|column_ref| { + row_slice[apc_poly_id_to_index_dense + [column_ref.id as usize]] + }) + } + } + }; + } + + // Evaluate bus interactions (same as standard path) + for ci in &compiled_interactions { + let mult = ci.mult.eval(row_slice); + periphery_real.apply( + ci.id as u16, + mult.as_canonical_u32(), + ci.args.iter().map(|a| a.eval(row_slice).as_canonical_u32()), + periphery_bus_ids, + ); + } }); - RowMajorMatrix::new(values, width) + if jit_debug { + tracing::info!("JIT_DEBUG: JIT trace gen complete: {}x{}", height, width); + if height > 0 { + let row0_u32: Vec = values[..width] + .iter() + .map(|v| v.as_canonical_u32()) + .collect(); + tracing::info!("JIT_DEBUG: JIT row0: {:?}", row0_u32); + } + } + + Ok(RowMajorMatrix::new(values, width)) + } + + /// Validation mode: runs the standard path, then independently computes + /// what the JIT path would produce for supported AIR types, and compares. + /// Returns the standard-path trace (so the proof is unaffected). + pub fn generate_witness_and_validate_jit( + &self, + original_arenas: OriginalArenas>, + ) -> DenseMatrix { + use std::collections::BTreeMap; + + let width = self.apc.machine().main_columns().count(); + + let mut original_arenas = match original_arenas { + OriginalArenas::Initialized(arenas) => arenas, + OriginalArenas::Uninitialized => { + return RowMajorMatrix::new(vec![], width); + } + }; + + let num_apc_calls = original_arenas.number_of_calls; + + // Build instruction metadata (same as in generate_witness_jit) + let instructions_with_subs: Vec<_> = self + .apc + .instructions() + .zip_eq(self.apc.subs().iter()) + .filter(|(_, subs)| !subs.is_empty()) + .collect(); + + let air_names: Vec = instructions_with_subs + .iter() + .map(|(instr, _)| { + let (air_id, _) = self.original_airs.get_instruction_air_and_id(instr); + air_id + }) + .collect(); + + let air_id_occurrences: HashMap<&str, usize> = + air_names.iter().map(|s| s.as_str()).counts(); + + let instruction_offsets: Vec = air_names + .iter() + .scan( + HashMap::<&str, usize>::default(), + |counts, air_name| { + let count = counts.entry(air_name.as_str()).or_default(); + let current = *count; + *count += 1; + Some(current) + }, + ) + .collect(); + + let apc_poly_id_to_index: BTreeMap = self + .apc + .machine() + .main_columns() + .enumerate() + .map(|(index, c)| (c.id, index)) + .collect(); + + let subs_by_instruction: Vec> = instructions_with_subs + .iter() + .map(|(_, subs)| { + subs.iter() + .map(|s| (s.original_poly_index, apc_poly_id_to_index[&s.apc_poly_id])) + .collect() + }) + .collect(); + + let mappings_by_air: HashMap<&str, jit_mapping::AirColumnMapping> = { + let mut m = HashMap::new(); + m.insert( + "VmAirWrapper", + jit_mapping::base_alu_mapping(), + ); + m.insert( + "VmAirWrapper", + jit_mapping::loadstore_mapping(), + ); + m.insert( + "VmAirWrapper", + jit_mapping::shift_mapping(), + ); + m.insert( + "VmAirWrapper", + jit_mapping::branch_equal_mapping(), + ); + m + }; + + let range_max_bits = self.periphery.real.range_checker.air.bus.range_max_bits as u32; + + // Clone arena bytes for the AIR types we can JIT, BEFORE the standard path consumes them. + // We use peek_real_arena to borrow without taking. + let mut jit_arena_bytes: HashMap, usize)> = HashMap::new(); + for air_name in air_names.iter().unique() { + if mappings_by_air.contains_key(air_name.as_str()) { + if let Some(arena) = original_arenas.peek_real_arena(air_name) { + let arena_width = arena.width; + let byte_vec: Vec = unsafe { + let ptr = arena.trace_buffer.as_ptr() as *const u8; + let len = arena.trace_buffer.len() * 4; + std::slice::from_raw_parts(ptr, len).to_vec() + }; + jit_arena_bytes.insert(air_name.clone(), (byte_vec, arena_width)); + } + } + } + + // Run standard path (consumes arenas) + let reference_trace = self + .generate_witness(OriginalArenas::Initialized(original_arenas)); + + // Now compute JIT values and compare with the reference trace + let mut mismatches = 0usize; + let mut checked = 0usize; + + for row_idx in 0..num_apc_calls { + let ref_row = &reference_trace.values[row_idx * width..(row_idx + 1) * width]; + + for (instr_idx, ((air_name, offset), subs)) in air_names + .iter() + .zip(instruction_offsets.iter()) + .zip(subs_by_instruction.iter()) + .enumerate() + { + let mapping = match mappings_by_air.get(air_name.as_str()) { + Some(m) => m, + None => continue, // Skip unsupported AIR types + }; + + let (arena_bytes, arena_width) = match jit_arena_bytes.get(air_name) { + Some(x) => (x.0.as_slice(), x.1), + None => continue, + }; + + let occurrences = *air_id_occurrences.get(air_name.as_str()).unwrap(); + let arena_row_idx = row_idx * occurrences + offset; + let row_bytes = arena_width * 4; + let record_byte_start = arena_row_idx * row_bytes; + let record_byte_end = record_byte_start + row_bytes; + + if record_byte_end > arena_bytes.len() { + continue; + } + let record_bytes = &arena_bytes[record_byte_start..record_byte_end]; + + for &(original_poly_index, apc_col_index) in subs { + let col_mapping = &mapping.columns[original_poly_index]; + let jit_val: BabyBear = jit_mapping::eval_column( + &col_mapping.computation, + record_bytes, + range_max_bits, + ); + let ref_val = ref_row[apc_col_index]; + checked += 1; + + if jit_val != ref_val { + mismatches += 1; + if mismatches <= 20 { + tracing::error!( + "JIT MISMATCH: row={}, instr={}, air='{}', orig_col={}, apc_col={}: JIT={}, REF={}", + row_idx, instr_idx, air_name, original_poly_index, apc_col_index, + jit_val.as_canonical_u32(), ref_val.as_canonical_u32() + ); + } + } + } + } + } + + if mismatches > 0 { + tracing::error!( + "JIT VALIDATION FAILED: {mismatches}/{checked} mismatches across {} rows", + num_apc_calls + ); + } else { + tracing::info!( + "JIT VALIDATION PASSED: {checked} column values checked across {} rows, all match!", + num_apc_calls + ); + } + + reference_trace } } diff --git a/openvm/src/powdr_extension/trace_generator/cpu/periphery.rs b/openvm/src/powdr_extension/trace_generator/cpu/periphery.rs index f88a964501..54b583ed5f 100644 --- a/openvm/src/powdr_extension/trace_generator/cpu/periphery.rs +++ b/openvm/src/powdr_extension/trace_generator/cpu/periphery.rs @@ -12,13 +12,13 @@ use openvm_circuit_primitives::{ range_tuple::{RangeTupleCheckerAir, RangeTupleCheckerChip, SharedRangeTupleCheckerChip}, var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerAir}, }; -use openvm_stark_backend::{config::StarkGenericConfig, p3_field::PrimeField32}; +use openvm_stark_backend::{p3_field::PrimeField32, StarkProtocolConfig}; use itertools::Itertools; use openvm_circuit::arch::RowMajorMatrixArena; -use openvm_stark_backend::config::Val; -use openvm_stark_backend::engine::StarkEngine; -use openvm_stark_backend::prover::cpu::{CpuBackend, CpuDevice}; +use openvm_cpu_backend::{CpuBackend, CpuDevice}; +use openvm_stark_backend::StarkEngine; +use openvm_stark_backend::Val; use crate::{isa::OpenVmISA, PeripheryBusIds}; @@ -77,7 +77,9 @@ impl PowdrPeripheryInstancesCpu { } } -impl VmExecutionExtension for SharedPeripheryChipsCpu { +impl, ISA: OpenVmISA> + VmExecutionExtension for SharedPeripheryChipsCpu +{ type Executor = DummyExecutor; fn extend_execution( @@ -89,7 +91,7 @@ impl VmExecutionExtension for SharedPeripher } } -impl VmCircuitExtension +impl VmCircuitExtension for SharedPeripheryChipsCpu { fn extend_circuit(&self, inventory: &mut AirInventory) -> Result<(), AirInventoryError> { @@ -132,7 +134,9 @@ pub struct SharedPeripheryChipsCpuProverExt; impl VmProverExtension> for SharedPeripheryChipsCpuProverExt where - SC: StarkGenericConfig, + SC: StarkProtocolConfig, + SC::F: Ord + openvm_stark_backend::p3_field::InjectiveMonomial<7>, + SC::EF: Ord, E: StarkEngine, PD = CpuDevice>, RA: RowMajorMatrixArena>, Val: PrimeField32, diff --git a/openvm/src/powdr_extension/trace_generator/cuda/inventory.rs b/openvm/src/powdr_extension/trace_generator/cuda/inventory.rs index 48a7d48c05..367c30aae4 100644 --- a/openvm/src/powdr_extension/trace_generator/cuda/inventory.rs +++ b/openvm/src/powdr_extension/trace_generator/cuda/inventory.rs @@ -2,7 +2,7 @@ use openvm_circuit::{ arch::{DenseRecordArena, VmChipComplex}, system::cuda::SystemChipInventoryGPU, }; -use openvm_cuda_backend::prover_backend::GpuBackend; +use openvm_cuda_backend::GpuBackend; pub type GpuDummyChipComplex = VmChipComplex; diff --git a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs index 6d439c3210..4bec236b22 100644 --- a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs @@ -5,12 +5,12 @@ use openvm_circuit::{ arch::{ChipInventory, DenseRecordArena}, utils::next_power_of_two_or_zero, }; +use openvm_circuit_primitives::Chip; use openvm_cuda_backend::base::DeviceMatrix; use openvm_cuda_common::copy::MemCopyH2D; use openvm_stark_backend::{ p3_field::PrimeField32, - prover::{hal::ProverBackend, types::AirProvingContext}, - Chip, + prover::{AirProvingContext, ProverBackend}, }; use openvm_stark_sdk::p3_baby_bear::BabyBear; use powdr_autoprecompiles::{ @@ -29,8 +29,223 @@ use crate::{ }; mod inventory; +pub mod nvrtc_bus_emit; +pub mod nvrtc_cache; +pub mod nvrtc_emit; mod periphery; +/// Selects the GPU trace-generation backend for `PowdrChipGpu`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum JitBackend { + /// Skip JIT entirely; use the original `apc_tracegen` substitution path. + Off, + /// Data-driven `apc_jit_tracegen` interpreter (current default JIT). + Descriptor, + /// NVRTC-compiled per-APC kernel, with per-APC fall-back to `Descriptor` + /// when the emitter does not yet support every column computation. + Nvrtc, +} + +/// Optional VPMM + first-H2D warm-up: pre-allocate a buffer of the requested +/// size, do a dummy H2D into it from a small pageable buffer, and drop. This +/// pays any first-use cost (VPMM `cuMemMap`, pinned staging pool init, +/// pinned-driver bookkeeping) before the next timed allocation/H2D. +fn maybe_warmup_vpmm(target_bytes: usize) { + if std::env::var("POWDR_VPMM_WARMUP").is_err() { + return; + } + let start = std::time::Instant::now(); + { + // Small probe to touch the small-allocation pool (cudaMallocAsync). + let _small = openvm_cuda_common::d_buffer::DeviceBuffer::::with_capacity(4096); + } + if target_bytes > 0 { + use openvm_cuda_common::copy::cuda_memcpy; + use std::ffi::c_void; + // Large probe to touch the VPMM (cuMemMap) path at the same scale we + // will allocate for real. + let large = openvm_cuda_common::d_buffer::DeviceBuffer::::with_capacity(target_bytes); + // Dummy H2D from a small pageable host buffer to pay any first-use + // pinned-staging cost. Copy size is small so the transfer itself is + // cheap; what we want is the warm-up of driver internal state. + // Probe sized to half the target (or 1 MB minimum), bounded at 256 MB + // so we don't double the runtime when warming up huge arenas. + let probe_size = (target_bytes / 2).max(1 << 20).min(256 << 20); + let probe: Vec = vec![0u8; probe_size]; + unsafe { + let _ = cuda_memcpy::( + large.as_mut_ptr() as *mut c_void, + probe.as_ptr() as *const c_void, + probe.len(), + ); + } + let _ = openvm_cuda_common::stream::current_stream_sync(); + drop(large); + } + tracing::info!( + "[trace_profile] vpmm_warmup {:8.3} ms (target {} bytes; H2D'd up to 256 MB probe)", + start.elapsed().as_secs_f64() * 1000.0, + target_bytes + ); +} + +/// Copy each arena's bytes directly to a single contiguous `DeviceBuffer` +/// without first concatenating on the host. Returns the buffer plus the +/// per-AIR byte offsets within it. The arenas Vec is kept alive until the +/// call returns so the H2D copies (issued async on the per-thread stream) +/// have valid source pointers — a stream sync at the next kernel launch +/// ensures completion before the kernel reads. +fn concat_arenas_direct_h2d( + iter: I, +) -> ( + openvm_cuda_common::d_buffer::DeviceBuffer, + std::collections::HashMap, +) +where + I: IntoIterator, +{ + use openvm_cuda_common::copy::cuda_memcpy; + use openvm_cuda_common::d_buffer::DeviceBuffer; + use std::ffi::c_void; + + let fn_start = std::time::Instant::now(); + let arenas: Vec<(String, openvm_circuit::arch::DenseRecordArena)> = iter + .into_iter() + .filter(|(_, a)| !a.allocated().is_empty()) + .collect(); + + let total_size: usize = arenas.iter().map(|(_, a)| a.allocated().len()).sum(); + if total_size == 0 { + return (DeviceBuffer::new(), std::collections::HashMap::new()); + } + let per_call_profile = std::env::var("POWDR_PER_CALL_PROFILE").is_ok(); + if per_call_profile { + tracing::info!( + "[per_call] concat.fn_entry_to_alloc {:8.3} ms", + fn_start.elapsed().as_secs_f64() * 1000.0 + ); + } + + let alloc_start = std::time::Instant::now(); + let dst = DeviceBuffer::::with_capacity(total_size); + if per_call_profile { + tracing::info!( + "[per_call] concat.alloc {:8.3} ms ({} bytes)", + alloc_start.elapsed().as_secs_f64() * 1000.0, + total_size + ); + } + + let mut offsets: std::collections::HashMap = std::collections::HashMap::new(); + let mut cursor: usize = 0; + for (name, arena) in &arenas { + let bytes = arena.allocated(); + let len = bytes.len(); + let call_start = std::time::Instant::now(); + // Safety: dst was allocated with total_size; cursor + len <= total_size by construction. + unsafe { + let dst_ptr = dst.as_mut_ptr().add(cursor) as *mut c_void; + let src_ptr = bytes.as_ptr() as *const c_void; + cuda_memcpy::(dst_ptr, src_ptr, len).expect("arena H2D failed"); + } + if per_call_profile { + let _ = openvm_cuda_common::stream::current_stream_sync(); + let dur_ms = call_start.elapsed().as_secs_f64() * 1000.0; + let bw = if dur_ms > 0.0 { + (len as f64) / 1.0e6 / dur_ms + } else { + 0.0 + }; + tracing::info!( + "[per_call] concat.memcpy[{:<60}] {:8.3} ms ({} bytes, {:.2} GB/s)", + name.chars().take(60).collect::(), + dur_ms, + len, + bw + ); + } + offsets.insert(name.clone(), cursor as u32); + cursor += len; + } + if per_call_profile { + // Time the arena Vec drop separately — for keccak APC=1 with 10k + // input this is ~65 ms of host-side free() on 1.65 GB of pageable + // memory. Without this measurement the cost hides inside + // time_stage's wrapper as if it were H2D time. + let pre_drop = fn_start.elapsed().as_secs_f64() * 1000.0; + let drop_start = std::time::Instant::now(); + std::mem::drop(arenas); + let drop_ms = drop_start.elapsed().as_secs_f64() * 1000.0; + tracing::info!( + "[per_call] concat.fn_pre_drop {:8.3} ms (H2D + alloc only)", + pre_drop + ); + tracing::info!( + "[per_call] concat.arenas_drop {:8.3} ms (host free of arena bytes)", + drop_ms + ); + tracing::info!( + "[per_call] concat.fn_total_to_return {:8.3} ms", + fn_start.elapsed().as_secs_f64() * 1000.0 + ); + return (dst, offsets); + } + // arenas Vec stays alive until function return. cudaMemcpyAsync source + // pointers are valid until the per-thread stream syncs, which it will at + // the next kernel launch / sync point in the caller. + let _ = arenas; // explicit reminder + (dst, offsets) +} + +/// Run `f`, optionally synchronizing the GPU stream and logging the elapsed +/// time. Active only when `POWDR_TRACE_PROFILE` is set in the environment; +/// otherwise this is a zero-overhead pass-through. +#[inline] +fn time_stage(name: &str, f: F) -> R +where + F: FnOnce() -> R, +{ + if std::env::var("POWDR_TRACE_PROFILE").is_err() { + return f(); + } + let start = std::time::Instant::now(); + let r = f(); + // Flush any GPU work this stage launched so the elapsed time reflects + // device-side execution rather than enqueue latency. + let _ = openvm_cuda_common::stream::current_stream_sync(); + let dur = start.elapsed(); + tracing::info!("[trace_profile] {:32} {:8.3} ms", name, dur.as_secs_f64() * 1000.0); + r +} + +/// Resolve the active JIT backend from the environment. +/// +/// `POWDR_JIT_BACKEND` takes precedence; values are case-insensitive +/// `off` / `descriptor` / `nvrtc`. For backwards compatibility, when +/// `POWDR_JIT_BACKEND` is unset the legacy `POWDR_JIT_TRACEGEN` flag enables +/// the `Descriptor` backend. +pub fn pick_jit_backend() -> JitBackend { + if let Ok(v) = std::env::var("POWDR_JIT_BACKEND") { + return match v.to_ascii_lowercase().as_str() { + "off" => JitBackend::Off, + "descriptor" => JitBackend::Descriptor, + "nvrtc" => JitBackend::Nvrtc, + other => { + tracing::warn!( + "Unknown POWDR_JIT_BACKEND='{}', defaulting to Descriptor", + other + ); + JitBackend::Descriptor + } + }; + } + if std::env::var("POWDR_JIT_TRACEGEN").is_ok() { + JitBackend::Descriptor + } else { + JitBackend::Off + } +} + pub use inventory::GpuDummyChipComplex; pub use periphery::{ PowdrPeripheryInstancesGpu, SharedPeripheryChipsGpu, SharedPeripheryChipsGpuProverExt, @@ -81,6 +296,53 @@ fn emit_expr( } /// Given the current bytecode, appends bytecode for the expression `expr` and returns the associated span +/// Walk an algebraic expression, inserting every `Reference.id` into `set`. +/// Used as a diagnostic: counts the distinct columns an expression touches, +/// which tells us if a "multi-column affine" decoder would cover it. +fn count_refs( + expr: &AlgebraicExpression, + set: &mut std::collections::HashSet, +) { + match expr { + AlgebraicExpression::Reference(r) => { + set.insert(r.id); + } + AlgebraicExpression::UnaryOperation(u) => count_refs(&u.expr, set), + AlgebraicExpression::BinaryOperation(b) => { + count_refs(&b.left, set); + count_refs(&b.right, set); + } + AlgebraicExpression::Number(_) => {} + } +} + +/// Structural fingerprint for an expression. Replaces specific ids and +/// constants with placeholders so identical SHAPES (different col indices +/// or constants) collapse to one summary string. Used to histogram the +/// unhandled bus-interaction shapes. +fn shape_summary(expr: &AlgebraicExpression) -> String { + match expr { + AlgebraicExpression::Number(_) => "N".to_string(), + AlgebraicExpression::Reference(_) => "R".to_string(), + AlgebraicExpression::UnaryOperation(u) => { + format!("neg({})", shape_summary(&u.expr)) + } + AlgebraicExpression::BinaryOperation(b) => { + let op = match b.op { + AlgebraicBinaryOperator::Add => "+", + AlgebraicBinaryOperator::Sub => "-", + AlgebraicBinaryOperator::Mul => "*", + }; + format!( + "({}{}{})", + shape_summary(&b.left), + op, + shape_summary(&b.right) + ) + } + } +} + fn emit_expr_span( bc: &mut Vec, expr: &AlgebraicExpression, @@ -176,11 +438,29 @@ pub fn compile_bus_to_gpu( (interactions, arg_spans, bytecode) } +/// Lazily-initialized per-APC codegen state. Holds the partition, the +/// per-APC source-key, and the four compiled kernel handles. Built on +/// first prove that uses `POWDR_BUS_CODEGEN=1`; reused for every +/// subsequent prove of the same APC. Eliminates the per-prove cost of +/// rebuilding source strings, hashing, and looking up the cache. +struct CachedBusCodegenKernels { + partition: nvrtc_bus_emit::PartitionedBus, + apc_poly_id_to_index: BTreeMap, + var_compiled: Option>, + tup_compiled: Option>, + br_compiled: Option>, + bx_compiled: Option>, +} + pub struct PowdrTraceGeneratorGpu { pub apc: IsaApc, pub original_airs: OriginalAirs, pub config: OriginalVmConfig, pub periphery: PowdrPeripheryInstancesGpu, + /// Per-APC codegen kernel cache; initialized on first call to + /// `launch_nvrtc_bus_codegen`. `OnceLock` makes the prove-path + /// kernel-fetch a single atomic load on warm runs. + codegen_cache: std::sync::OnceLock>, } impl PowdrTraceGeneratorGpu { @@ -190,11 +470,803 @@ impl PowdrTraceGeneratorGpu { config: OriginalVmConfig, periphery: PowdrPeripheryInstancesGpu, ) -> Self { - Self { + let gen = Self { apc, original_airs, config, periphery, + codegen_cache: std::sync::OnceLock::new(), + }; + + // When NVRTC codegen is selected via env, eagerly compile the + // bus kernels at construction time so the cost moves to process + // startup (or APC-load time) and doesn't hit any prove timer. + // The kernel sources are APC-deterministic and apc_height- + // independent, so this is safe to do once per APC. + if std::env::var("POWDR_JIT_BACKEND_BUS") + .map(|v| v == "nvrtc") + .unwrap_or(false) + && std::env::var("POWDR_BUS_CODEGEN").is_ok() + { + gen.warm_codegen_cache(); + } + + gen + } + + /// Eagerly populate `codegen_cache`. Call from APC construction or + /// process-startup paths to amortize NVRTC compile out of the prove + /// timer. Idempotent — `OnceLock::get_or_init` ensures it runs only + /// the first time. If NVRTC codegen isn't configured at runtime, + /// callers can skip this without affecting correctness. + pub fn warm_codegen_cache(&self) { + let apc_poly_id_to_index: BTreeMap = self + .apc + .machine + .main_columns() + .enumerate() + .map(|(index, c)| (c.id, index)) + .collect(); + + let var_range_bus_id = self.periphery.bus_ids.range_checker as u32; + let tuple2_bus_id = self + .periphery + .bus_ids + .tuple_range_checker + .map(|x| x as u32) + .unwrap_or(0); + let bitwise_bus_id = self + .periphery + .bus_ids + .bitwise_lookup + .map(|x| x as u32) + .unwrap_or(0); + + // height parameter is unused by codegen (kernel takes H as runtime + // arg); pick any value that satisfies the partition's interface. + let _ = self.codegen_cache.get_or_init(|| { + time_stage("bus_kernel_nvrtc_codegen.warm_cache", || { + self.build_codegen_cache( + &apc_poly_id_to_index, + 1, // dummy height — not baked into source + var_range_bus_id, + tuple2_bus_id, + bitwise_bus_id, + ) + }) + }); + } + + /// Run the periphery-bus pass on a populated trace. Identical work in + /// all three trace-gen paths (baseline / descriptor / nvrtc) — exists + /// to dedupe what was a 3× copy-paste. + /// + /// `stage_prefix` controls the [trace_profile] labels for sub-stage + /// timing (e.g., `"baseline"`, `"descriptor"`, `"nvrtc"`). + /// + /// Backend selection (env-driven; first-match wins): + /// - `POWDR_BUS_SPIKE` → run the spike kernel (Phase 0 measurement; NOT + /// a real proof — histograms are MEANINGLESS). + /// - `POWDR_JIT_BACKEND_BUS=nvrtc` → NVRTC bus kernel. + /// - default → existing bytecode-VM kernel. + fn run_bus_pass( + &self, + output: &DeviceMatrix, + num_apc_calls: usize, + apc_poly_id_to_index: &BTreeMap, + height: usize, + stage_prefix: &str, + ) { + let backend_nvrtc = std::env::var("POWDR_JIT_BACKEND_BUS") + .map(|v| v == "nvrtc") + .unwrap_or(false); + + // The bytecode-VM kernel and the spike kernel both consume the + // pre-compiled bytecode. The NVRTC kernel does not — but we still + // pay the H2D for compatibility (cheap; ~10us). Skip when NVRTC. + let compile_h2d = !backend_nvrtc || std::env::var("POWDR_BUS_SPIKE").is_ok(); + let bus_data = if compile_h2d { + time_stage(&format!("{stage_prefix}.bus_compile_h2d"), || { + let (bus_interactions, arg_spans, bytecode) = compile_bus_to_gpu( + &self.apc.machine.bus_interactions, + apc_poly_id_to_index, + height, + ); + let bus_interactions = bus_interactions.to_device().unwrap(); + let arg_spans = arg_spans.to_device().unwrap(); + let bytecode = bytecode.to_device().unwrap(); + Some((bus_interactions, arg_spans, bytecode)) + }) + } else { + None + }; + + let periphery = &self.periphery.real; + let var_range_bus_id = self.periphery.bus_ids.range_checker as u32; + let var_range_count = &periphery.range_checker.count; + let tuple_range_checker_chip = periphery.tuple_range_checker.as_ref().unwrap(); + let tuple2_bus_id = self.periphery.bus_ids.tuple_range_checker.unwrap() as u32; + let tuple2_sizes = tuple_range_checker_chip.sizes; + let tuple2_count_u32 = tuple_range_checker_chip.count.as_ref(); + let bitwise_bus_id = self.periphery.bus_ids.bitwise_lookup.unwrap() as u32; + let bitwise_count_u32 = periphery.bitwise_lookup_8.as_ref().unwrap().count.as_ref(); + + if std::env::var("POWDR_BUS_SPIKE").is_ok() { + let (bus_interactions, _, _) = bus_data.expect("spike requires H2D"); + time_stage(&format!("{stage_prefix}.bus_kernel_spike"), || { + cuda_abi::apc_apply_bus_spike( + output, + num_apc_calls, + bus_interactions, + var_range_bus_id, + var_range_count, + tuple2_bus_id, + tuple2_count_u32, + tuple2_sizes, + bitwise_bus_id, + bitwise_count_u32, + ) + .unwrap(); + }); + return; + } + + if backend_nvrtc { + let codegen = std::env::var("POWDR_BUS_CODEGEN").is_ok(); + let stage_label = if codegen { + format!("{stage_prefix}.bus_kernel_nvrtc_codegen") + } else { + format!("{stage_prefix}.bus_kernel_nvrtc") + }; + time_stage(&stage_label, || { + if codegen { + self.launch_nvrtc_bus_codegen( + output, + num_apc_calls, + apc_poly_id_to_index, + height, + var_range_bus_id, + var_range_count, + tuple2_bus_id, + tuple2_count_u32, + tuple2_sizes, + bitwise_bus_id, + bitwise_count_u32, + ); + } else { + self.launch_nvrtc_bus( + output, + num_apc_calls, + apc_poly_id_to_index, + height, + var_range_bus_id, + var_range_count, + tuple2_bus_id, + tuple2_count_u32, + tuple2_sizes, + bitwise_bus_id, + bitwise_count_u32, + ); + } + }); + return; + } + + let (bus_interactions, arg_spans, bytecode) = bus_data.expect("vm requires H2D"); + time_stage(&format!("{stage_prefix}.bus_kernel"), || { + 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(); + }); + } + + /// Per-APC codegen path: emit 4 kernels (one per bus-kind that has + /// captured ops) with op constants baked as immediates in source. Each + /// kernel's source is unique per APC's op set; the PTX disk cache + /// dedupes when partitions match across runs/APCs. + #[allow(clippy::too_many_arguments)] + /// Eager build of the per-APC codegen kernels: partition + codegen + + /// NVRTC compile + module load. Called once per APC on first prove + /// (lazy via `codegen_cache.get_or_init`). The result is cached for + /// every subsequent prove of this APC, so no source rebuild / cache + /// lookup happens on the hot prove path. + /// + /// Note: `height` (= `next_power_of_two_or_zero(num_apc_calls)`) is + /// per-prove and intentionally NOT baked into the kernel source — + /// the kernel takes it as the runtime arg `H`. The cache is therefore + /// height-independent and can be reused across proves of any size. + fn build_codegen_cache( + &self, + apc_poly_id_to_index: &BTreeMap, + height: usize, + var_range_bus_id: u32, + tuple2_bus_id: u32, + bitwise_bus_id: u32, + ) -> std::sync::Arc { + use crate::powdr_extension::trace_generator::cuda::nvrtc_bus_emit::{ + build_emitter_input, emit_codegen_bitwise, emit_codegen_tuple2, + emit_codegen_var_range, partition_apc_bus, + }; + use crate::powdr_extension::trace_generator::cuda::nvrtc_cache::NvrtcKernelCache; + + let p = time_stage("bus_kernel_nvrtc_codegen.cache_build.partition", || { + let input = build_emitter_input( + &self.apc.machine.bus_interactions, + apc_poly_id_to_index, + height, + var_range_bus_id as u64, + Some(tuple2_bus_id as u64), + Some(bitwise_bus_id as u64), + ); + partition_apc_bus(&input) + .unwrap_or_else(|e| panic!("NVRTC bus partition failed: {}", e)) + }); + + let kernels = time_stage("bus_kernel_nvrtc_codegen.cache_build.codegen", || { + ( + emit_codegen_var_range(&p.var_ops, &p.var_ops_bilinear), + emit_codegen_tuple2(&p.tuple_ops, &p.tuple_ops_bilinear), + emit_codegen_bitwise( + &p.bitwise_range_ops, + &p.bitwise_range_ops_bilinear, + false, + ), + emit_codegen_bitwise( + &p.bitwise_xor_ops, + &p.bitwise_xor_ops_bilinear, + true, + ), + ) + }); + + // Once-per-process source-size diagnostic (only the first APC's). + { + use std::sync::atomic::{AtomicBool, Ordering}; + static REPORTED: AtomicBool = AtomicBool::new(false); + if !REPORTED.swap(true, Ordering::Relaxed) { + let bytes_for = + |k: &Option| { + k.as_ref().map(|x| x.source.len()).unwrap_or(0) + }; + tracing::info!( + "NVRTC bus codegen (first APC, build-time): src bytes var={}, tup={}, br={}, bx={}, total={}", + bytes_for(&kernels.0), + bytes_for(&kernels.1), + bytes_for(&kernels.2), + bytes_for(&kernels.3), + bytes_for(&kernels.0) + + bytes_for(&kernels.1) + + bytes_for(&kernels.2) + + bytes_for(&kernels.3) + ); + } + } + + let compiled = time_stage("bus_kernel_nvrtc_codegen.cache_build.compile", || { + let kernels_vec: Vec<_> = [&kernels.0, &kernels.1, &kernels.2, &kernels.3] + .iter() + .filter_map(|k| (*k).clone()) + .collect(); + let names: Vec<_> = kernels_vec.iter().map(|k| k.name.clone()).collect(); + let compiled = NvrtcKernelCache::global().get_or_compile_many(&kernels_vec); + compiled + .into_iter() + .enumerate() + .map(|(i, r)| { + r.unwrap_or_else(|e| { + panic!("NVRTC compile failed for {}: {:?}", names[i], e) + }) + }) + .collect::>() + }); + + // Map filtered-compiled list back to per-kind Optional handles. + let mut compiled_iter = compiled.into_iter(); + let var_compiled = kernels.0.as_ref().map(|_| compiled_iter.next().unwrap()); + let tup_compiled = kernels.1.as_ref().map(|_| compiled_iter.next().unwrap()); + let br_compiled = kernels.2.as_ref().map(|_| compiled_iter.next().unwrap()); + let bx_compiled = kernels.3.as_ref().map(|_| compiled_iter.next().unwrap()); + + std::sync::Arc::new(CachedBusCodegenKernels { + partition: p, + apc_poly_id_to_index: apc_poly_id_to_index.clone(), + var_compiled, + tup_compiled, + br_compiled, + bx_compiled, + }) + } + + fn launch_nvrtc_bus_codegen( + &self, + output: &DeviceMatrix, + num_apc_calls: usize, + apc_poly_id_to_index: &BTreeMap, + height: usize, + var_range_bus_id: u32, + var_range_count: &openvm_cuda_common::d_buffer::DeviceBuffer, + tuple2_bus_id: u32, + tuple2_count: &openvm_cuda_common::d_buffer::DeviceBuffer, + tuple2_sizes: [u32; 2], + bitwise_bus_id: u32, + bitwise_count: &openvm_cuda_common::d_buffer::DeviceBuffer, + ) { + // Initialize the per-APC kernel cache on first prove. Subsequent + // proves of the same APC skip partition + codegen + compile entirely + // — they just clone the Arc handles via OnceLock. + let cached = time_stage("bus_kernel_nvrtc_codegen.cache_get", || { + self.codegen_cache.get_or_init(|| { + self.build_codegen_cache( + apc_poly_id_to_index, + height, + var_range_bus_id, + tuple2_bus_id, + bitwise_bus_id, + ) + }) + .clone() + }); + + let p = &cached.partition; + let var_compiled = &cached.var_compiled; + let tup_compiled = &cached.tup_compiled; + let br_compiled = &cached.br_compiled; + let bx_compiled = &cached.bx_compiled; + + // Launch each kind. Grid sized for n_ops; block 256 = 8 warps. + let block_x = 256u32; + let warps_per_block = 8u32; + + if let (Some(comp), n) = ( + var_compiled.as_ref(), + (p.var_ops.len() + p.var_ops_bilinear.len()) as u32, + ) { + time_stage("bus_kernel_nvrtc_codegen.launch_var", || { + let grid_x = n.div_ceil(warps_per_block).max(1); + let rc = unsafe { + cuda_abi::powdr_nvrtc_launch_bus_v4( + comp.function(), + output.buffer().as_ptr() as *const u32, + num_apc_calls as i32, + height as u64, + var_range_count.as_mut_ptr() as *mut u32, + var_range_count.len() as u32, + 0, + 0, + grid_x, + block_x, + ) + }; + assert_eq!(rc, 0, "var_range codegen launch rc={}", rc); + }); + } + + if let (Some(comp), n) = ( + tup_compiled.as_ref(), + (p.tuple_ops.len() + p.tuple_ops_bilinear.len()) as u32, + ) { + time_stage("bus_kernel_nvrtc_codegen.launch_tup", || { + let grid_x = n.div_ceil(warps_per_block).max(1); + let rc = unsafe { + cuda_abi::powdr_nvrtc_launch_bus_v4( + comp.function(), + output.buffer().as_ptr() as *const u32, + num_apc_calls as i32, + height as u64, + tuple2_count.as_mut_ptr() as *mut u32, + tuple2_sizes[0], + tuple2_sizes[1], + 1, + grid_x, + block_x, + ) + }; + assert_eq!(rc, 0, "tuple2 codegen launch rc={}", rc); + }); + } + + if let (Some(comp), n) = ( + br_compiled.as_ref(), + (p.bitwise_range_ops.len() + p.bitwise_range_ops_bilinear.len()) as u32, + ) { + time_stage("bus_kernel_nvrtc_codegen.launch_br", || { + let grid_x = n.div_ceil(warps_per_block).max(1); + let rc = unsafe { + cuda_abi::powdr_nvrtc_launch_bus_v4_bitwise( + comp.function(), + output.buffer().as_ptr() as *const u32, + num_apc_calls as i32, + height as u64, + bitwise_count.as_mut_ptr() as *mut u32, + grid_x, + block_x, + ) + }; + assert_eq!(rc, 0, "bitwise_range codegen launch rc={}", rc); + }); + } + + if let (Some(comp), n) = ( + bx_compiled.as_ref(), + (p.bitwise_xor_ops.len() + p.bitwise_xor_ops_bilinear.len()) as u32, + ) { + time_stage("bus_kernel_nvrtc_codegen.launch_bx", || { + let grid_x = n.div_ceil(warps_per_block).max(1); + let rc = unsafe { + cuda_abi::powdr_nvrtc_launch_bus_v4_bitwise( + comp.function(), + output.buffer().as_ptr() as *const u32, + num_apc_calls as i32, + height as u64, + bitwise_count.as_mut_ptr() as *mut u32, + grid_x, + block_x, + ) + }; + assert_eq!(rc, 0, "bitwise_xor codegen launch rc={}", rc); + }); + } + + // Bytecode-VM fallback for unhandled (same as interpreter path). + if !p.unhandled.is_empty() { + let (bus_interactions, arg_spans, bytecode) = + time_stage("bus_kernel_nvrtc_codegen.fallback_compile_h2d", || { + let unhandled_interactions: Vec<_> = p + .unhandled + .iter() + .map(|&i| self.apc.machine.bus_interactions[i].clone()) + .collect(); + let (bus_interactions, arg_spans, bytecode) = compile_bus_to_gpu( + &unhandled_interactions, + apc_poly_id_to_index, + height, + ); + let bus_interactions = bus_interactions.to_device().unwrap(); + let arg_spans = arg_spans.to_device().unwrap(); + let bytecode = bytecode.to_device().unwrap(); + (bus_interactions, arg_spans, bytecode) + }); + time_stage("bus_kernel_nvrtc_codegen.fallback_vm", || { + 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, + tuple2_sizes, + bitwise_bus_id, + bitwise_count, + ) + .unwrap(); + }); + } + } + + /// Partition + compile + launch the NVRTC bus kernels for this APC. + /// Four kind-templated kernels (var_range, tuple2, bitwise_range, + /// bitwise_xor) share a single PTX each across all APCs and runs. + /// Per-APC op tables flow through `__constant__` memory uploaded via + /// `cuMemcpyHtoD` between launches. + #[allow(clippy::too_many_arguments)] + fn launch_nvrtc_bus( + &self, + output: &DeviceMatrix, + num_apc_calls: usize, + apc_poly_id_to_index: &BTreeMap, + height: usize, + var_range_bus_id: u32, + var_range_count: &openvm_cuda_common::d_buffer::DeviceBuffer, + tuple2_bus_id: u32, + tuple2_count: &openvm_cuda_common::d_buffer::DeviceBuffer, + tuple2_sizes: [u32; 2], + bitwise_bus_id: u32, + bitwise_count: &openvm_cuda_common::d_buffer::DeviceBuffer, + ) { + use crate::powdr_extension::trace_generator::cuda::nvrtc_bus_emit::{ + build_emitter_input, kernel_sources, partition_apc_bus, + }; + use crate::powdr_extension::trace_generator::cuda::nvrtc_cache::NvrtcKernelCache; + use std::ffi::CString; + + let input = build_emitter_input( + &self.apc.machine.bus_interactions, + apc_poly_id_to_index, + height, + var_range_bus_id as u64, + Some(tuple2_bus_id as u64), + Some(bitwise_bus_id as u64), + ); + + let p = time_stage("bus_kernel_nvrtc.partition", || { + partition_apc_bus(&input) + .unwrap_or_else(|e| panic!("NVRTC bus partition failed: {}", e)) + }); + + // Once-per-process diagnostic: report partition coverage and dump + // every unhandled interaction's shape + bytecode-cost proxy. The + // bytecode-cost weighting tells us which unhandled shapes account + // for most of the fallback_vm time; the shape histogram tells us + // whether multi-column affine support would catch them. + { + use std::collections::HashMap; + use std::sync::atomic::{AtomicBool, Ordering}; + static REPORTED: AtomicBool = AtomicBool::new(false); + if !REPORTED.swap(true, Ordering::Relaxed) { + let total = self.apc.machine.bus_interactions.len(); + let simple = p.var_ops.len() + + p.tuple_ops.len() + + p.bitwise_range_ops.len() + + p.bitwise_xor_ops.len(); + tracing::info!( + "NVRTC bus partition (first APC): total={}, simple={} (var={}, tup={}, br={}, bx={}), unsupported={}, unhandled={}", + total, simple, + p.var_ops.len(), p.tuple_ops.len(), + p.bitwise_range_ops.len(), p.bitwise_xor_ops.len(), + p.n_unsupported, p.unhandled.len() + ); + + // Per-unhandled stats. bytecode_size is a cost proxy: per-row + // VM cost is roughly proportional to bytecode_size since the + // VM walks each opcode in sequence. + let mut by_shape: HashMap = HashMap::new(); // shape → (count, total_bc) + let mut by_max_cols: HashMap = HashMap::new(); + let mut total_unhandled_bc: usize = 0; + + tracing::info!(" per-unhandled (sorted by bytecode_size desc):"); + let mut rows: Vec<_> = p + .unhandled + .iter() + .map(|&i| { + let bi = &self.apc.machine.bus_interactions[i]; + let mut bc: Vec = Vec::new(); + let _ = emit_expr_span(&mut bc, &bi.mult, apc_poly_id_to_index, height); + let mult_bc = bc.len(); + let mut arg_bc_total = 0usize; + let mut max_cols = 0usize; + let mut shape_summaries: Vec = Vec::new(); + for a in &bi.args { + bc.clear(); + let _ = emit_expr_span(&mut bc, a, apc_poly_id_to_index, height); + arg_bc_total += bc.len(); + let mut col_set = std::collections::HashSet::new(); + count_refs(a, &mut col_set); + max_cols = max_cols.max(col_set.len()); + shape_summaries.push(shape_summary(a)); + } + let total_bc = mult_bc + arg_bc_total; + let shape = shape_summaries.join(" | "); + (i, bi.id, total_bc, max_cols, shape) + }) + .collect(); + rows.sort_by(|a, b| b.2.cmp(&a.2)); + + for (i, bus_id, total_bc, max_cols, shape) in &rows { + total_unhandled_bc += total_bc; + let entry = by_shape + .entry(shape.clone()) + .or_insert((0usize, 0usize)); + entry.0 += 1; + entry.1 += total_bc; + *by_max_cols.entry(*max_cols).or_insert(0) += 1; + tracing::info!( + " idx={:5} bus_id={} bc={:3} max_cols={} shape={}", + i, bus_id, total_bc, max_cols, shape + ); + } + tracing::info!( + " total unhandled bytecode words: {}", + total_unhandled_bc + ); + tracing::info!(" unhandled bucketed by max_cols_in_any_arg:"); + let mut buckets: Vec<_> = by_max_cols.into_iter().collect(); + buckets.sort(); + for (cols, n) in &buckets { + tracing::info!(" {} col(s): {} interactions", cols, n); + } + tracing::info!(" unhandled bucketed by structural shape (top 10 by bytecode-weight):"); + let mut shapes: Vec<_> = by_shape.into_iter().collect(); + shapes.sort_by(|a, b| b.1 .1.cmp(&a.1 .1)); + for (shape, (cnt, bc)) in shapes.iter().take(10) { + tracing::info!(" n={:3} bc_total={:5} {}", cnt, bc, shape); + } + } + } + + // Prefetch all 4 compiled kernels in one batch (parallel-compiles + // cold, dedupes warm). Stage covers PTX cache hit + module load. + let compiled = time_stage("bus_kernel_nvrtc.kernel_prefetch", || { + let sources = kernel_sources(); + let kernels: Vec<_> = sources.iter().map(|(_, k)| k.clone()).collect(); + let names: Vec<_> = sources.iter().map(|(_, k)| k.name.clone()).collect(); + let compiled = NvrtcKernelCache::global().get_or_compile_many(&kernels); + compiled + .into_iter() + .enumerate() + .map(|(i, r)| { + r.unwrap_or_else(|e| { + panic!("NVRTC compile failed for kernel {}: {:?}", names[i], e) + }) + }) + .collect::>() + }); + + // Per-kind upload + launch with explicit timers. Each kind allocates + // a fresh DeviceBuffer, H2Ds the op slice into it, then launches + // with bus_v3 (kernel reads ops from the d_ops pointer arg). + // Ownership: bufs[i] kept alive until cuCtxSynchronize at end of + // launch, then dropped (each Drop triggers cudaFree). + let do_kind = |stage_upload: &str, + stage_launch: &str, + k_idx: usize, + ops_bytes: &[u8], + n_ops: u32, + d_hist: *mut u32, + extra0: u32, + extra1: Option| + -> Option> { + if n_ops == 0 { + return None; + } + // H2D the op table. + let d_ops_buf = time_stage(stage_upload, || { + use openvm_cuda_common::copy::MemCopyH2D; + ops_bytes.to_device().expect("op table H2D") + }); + // Launch. + time_stage(stage_launch, || { + let func = compiled[k_idx].function(); + let warps_per_block = 8u32; + let grid_x = n_ops.div_ceil(warps_per_block).max(1); + let block_x = 256u32; + let rc = unsafe { + cuda_abi::powdr_nvrtc_launch_bus_v3( + func, + output.buffer().as_ptr() as *const u32, + num_apc_calls as i32, + height as u64, + d_hist, + extra0, + extra1.unwrap_or(0), + if extra1.is_some() { 1 } else { 0 }, + n_ops, + d_ops_buf.as_ptr() as *const std::ffi::c_void, + grid_x, + block_x, + ) + }; + assert_eq!(rc, 0, "launch_bus_v3 rc={}", rc); + }); + Some(d_ops_buf) + }; + + // Keep buffers alive until end of scope (after all launches). + let _ = CString::new(""); // suppress unused-import warning if any + let var_bytes: &[u8] = unsafe { + std::slice::from_raw_parts( + p.var_ops.as_ptr() as *const u8, + std::mem::size_of_val(&p.var_ops[..]), + ) + }; + let _b0 = do_kind( + "bus_kernel_nvrtc.upload_var", + "bus_kernel_nvrtc.launch_var", + 0, + var_bytes, + p.var_ops.len() as u32, + var_range_count.as_mut_ptr() as *mut u32, + var_range_count.len() as u32, + None, + ); + + let tup_bytes: &[u8] = unsafe { + std::slice::from_raw_parts( + p.tuple_ops.as_ptr() as *const u8, + std::mem::size_of_val(&p.tuple_ops[..]), + ) + }; + let _b1 = do_kind( + "bus_kernel_nvrtc.upload_tup", + "bus_kernel_nvrtc.launch_tup", + 1, + tup_bytes, + p.tuple_ops.len() as u32, + tuple2_count.as_mut_ptr() as *mut u32, + tuple2_sizes[0], + Some(tuple2_sizes[1]), + ); + + let br_bytes: &[u8] = unsafe { + std::slice::from_raw_parts( + p.bitwise_range_ops.as_ptr() as *const u8, + std::mem::size_of_val(&p.bitwise_range_ops[..]), + ) + }; + let _b2 = do_kind( + "bus_kernel_nvrtc.upload_br", + "bus_kernel_nvrtc.launch_br", + 2, + br_bytes, + p.bitwise_range_ops.len() as u32, + bitwise_count.as_mut_ptr() as *mut u32, + 0, + None, + ); + + let bx_bytes: &[u8] = unsafe { + std::slice::from_raw_parts( + p.bitwise_xor_ops.as_ptr() as *const u8, + std::mem::size_of_val(&p.bitwise_xor_ops[..]), + ) + }; + let _b3 = do_kind( + "bus_kernel_nvrtc.upload_bx", + "bus_kernel_nvrtc.launch_bx", + 3, + bx_bytes, + p.bitwise_xor_ops.len() as u32, + bitwise_count.as_mut_ptr() as *mut u32, + 0, + None, + ); + + // Bytecode-VM fallback for the `unhandled` tail. The interpreter + // path doesn't support bilinear arg shapes — bilinear ops from the + // partition fall to VM here too. + let mut all_unhandled: Vec = p.unhandled.clone(); + all_unhandled.extend(p.var_ops_bilinear.iter().map(|(i, _)| *i)); + all_unhandled.extend(p.tuple_ops_bilinear.iter().map(|(i, _)| *i)); + all_unhandled.extend(p.bitwise_range_ops_bilinear.iter().map(|(i, _)| *i)); + all_unhandled.extend(p.bitwise_xor_ops_bilinear.iter().map(|(i, _)| *i)); + if !all_unhandled.is_empty() { + let (bus_interactions, arg_spans, bytecode) = + time_stage("bus_kernel_nvrtc.fallback_compile_h2d", || { + let unhandled_interactions: Vec<_> = all_unhandled + .iter() + .map(|&i| self.apc.machine.bus_interactions[i].clone()) + .collect(); + let (bus_interactions, arg_spans, bytecode) = compile_bus_to_gpu( + &unhandled_interactions, + apc_poly_id_to_index, + height, + ); + let bus_interactions = bus_interactions.to_device().unwrap(); + let arg_spans = arg_spans.to_device().unwrap(); + let bytecode = bytecode.to_device().unwrap(); + (bus_interactions, arg_spans, bytecode) + }); + time_stage("bus_kernel_nvrtc.fallback_vm", || { + 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, + tuple2_sizes, + bitwise_bus_id, + bitwise_count, + ) + .unwrap(); + }); } } @@ -212,41 +1284,116 @@ impl PowdrTraceGeneratorGpu { let num_apc_calls = original_arenas.number_of_calls; - let chip_inventory: ChipInventory = { - let airs = ISA::create_dummy_airs(self.config.config(), self.periphery.dummy.clone()) - .expect("Failed to create dummy airs"); + let chip_inventory: ChipInventory = time_stage( + "baseline.dummy_chip_inventory", + || { + let airs = ISA::create_dummy_airs(self.config.config(), self.periphery.dummy.clone()) + .expect("Failed to create dummy airs"); - ISA::create_dummy_chip_complex_gpu( - self.config.config(), - airs, - self.periphery.dummy.clone(), - ) - .expect("Failed to create chip complex") - .inventory - }; - - let dummy_trace_by_air_name: HashMap> = chip_inventory - .chips() - .iter() - .enumerate() - .rev() - .filter_map(|(insertion_idx, chip)| { - let air_name = chip_inventory.airs().ext_airs()[insertion_idx].name(); - - let record_arena = { - match original_arenas.take_real_arena(&air_name) { - Some(ra) => ra, - None => return None, // skip this iteration, because we only have record arena for chips that are used - } - }; + ISA::create_dummy_chip_complex_gpu( + self.config.config(), + airs, + self.periphery.dummy.clone(), + ) + .expect("Failed to create chip complex") + .inventory + }, + ); - // We might have initialized an arena for an AIR which ends up having no real records. It gets filtered out here. - chip.generate_proving_ctx(record_arena) - .common_main - .map(|m| (air_name, m)) - }) + // The set of AIR names that the APC's instructions reference. JIT paths + // bypass primitive `fill_trace_row` for these; baseline runs them. + // Tagging per-chip times with this set makes the APC-AIR sum here + // directly comparable to the JIT path's `arena_direct_h2d`. + let apc_air_names: std::collections::HashSet = self + .apc + .instructions() + .map(|instr| self.original_airs.opcode_to_air[&instr.inner.opcode].clone()) .collect(); + let mut apc_air_total_ms: f64 = 0.0; + let mut non_apc_air_total_ms: f64 = 0.0; + let mut apc_air_total_bytes: usize = 0; + let mut non_apc_air_total_bytes: usize = 0; + let profile_on = std::env::var("POWDR_TRACE_PROFILE").is_ok(); + + // Pre-take all arenas so we can do an isolated H2D-only measurement + // before consuming them via chip.generate_proving_ctx. + let chip_arena_triples: Vec<(usize, String, openvm_circuit::arch::DenseRecordArena)> = + chip_inventory + .chips() + .iter() + .enumerate() + .rev() + .filter_map(|(insertion_idx, _chip)| { + let air_name = chip_inventory.airs().ext_airs()[insertion_idx].name(); + original_arenas + .take_real_arena(&air_name) + .map(|a| (insertion_idx, air_name, a)) + }) + .collect(); + + let dummy_trace_by_air_name: HashMap> = time_stage( + "baseline.fill_dummy_traces", + || { + let map: HashMap<_, _> = chip_arena_triples + .into_iter() + .filter_map(|(insertion_idx, air_name, record_arena)| { + let chip = &chip_inventory.chips()[insertion_idx]; + let bytes_len = record_arena.allocated().len(); + let is_apc = apc_air_names.contains(&air_name); + + // Time each chip's generate_proving_ctx individually + // (covers H2D + fill_trace_row + alloc). + let ctx = if profile_on { + let start = std::time::Instant::now(); + let r = chip.generate_proving_ctx(record_arena); + let _ = openvm_cuda_common::stream::current_stream_sync(); + let ms = start.elapsed().as_secs_f64() * 1000.0; + let tag = if is_apc { "APC" } else { "non-APC" }; + tracing::info!( + "[trace_profile] baseline.fill_dummy_trace[{}/{:<70}] {:8.3} ms ({} bytes)", + tag, + air_name.chars().take(70).collect::(), + ms, + bytes_len + ); + if is_apc { + apc_air_total_ms += ms; + apc_air_total_bytes += bytes_len; + } else { + non_apc_air_total_ms += ms; + non_apc_air_total_bytes += bytes_len; + } + r + } else { + chip.generate_proving_ctx(record_arena) + }; + + let m = ctx.common_main; + use openvm_stark_backend::prover::MatrixDimensions; + if m.height() > 0 { + Some((air_name, m)) + } else { + None + } + }) + .collect(); + if profile_on { + tracing::info!( + "[trace_profile] baseline.fill_dummy_trace_APC_total {:8.3} ms ({} bytes)", + apc_air_total_ms, + apc_air_total_bytes + ); + tracing::info!( + "[trace_profile] baseline.fill_dummy_trace_nonAPC_total {:8.3} ms ({} bytes)", + non_apc_air_total_ms, + non_apc_air_total_bytes + ); + } + map + }, + ); + // Map from apc poly id to its index in the final apc trace let apc_poly_id_to_index: BTreeMap = self .apc @@ -256,13 +1403,18 @@ impl PowdrTraceGeneratorGpu { .map(|(index, c)| (c.id, index)) .collect(); - // allocate for apc trace + // allocate for apc trace (zero-initialized so columns not covered + // by substitutions or derived expressions default to zero, matching the CPU path) let width = apc_poly_id_to_index.len(); let height = next_power_of_two_or_zero(num_apc_calls); - let mut output = DeviceMatrix::::with_capacity(height, width); + let mut output = time_stage("baseline.allocate_zero_output", || { + let m = DeviceMatrix::::with_capacity(height, width); + m.buffer().fill_zero().unwrap(); + m + }); // Prepare `OriginalAir` and `Subst` arrays - let (airs, substitutions) = { + let (airs, substitutions) = time_stage("baseline.build_subs", || { self.apc // go through original instructions .instructions() @@ -305,7 +1457,7 @@ impl PowdrTraceGeneratorGpu { // get the device dummy trace for this air let dummy_trace = &dummy_trace_by_air_name[*air_name]; - use openvm_stark_backend::prover::hal::MatrixDimensions; + use openvm_stark_backend::prover::MatrixDimensions; airs.push(OriginalAir { width: dummy_trace.width() as i32, height: dummy_trace.height() as i32, @@ -318,80 +1470,1075 @@ impl PowdrTraceGeneratorGpu { (airs, substitutions) }, ) - }; + }); // Send the airs and substitutions to device - let airs = airs.to_device().unwrap(); - let substitutions = substitutions.to_device().unwrap(); + let (airs, substitutions) = time_stage("baseline.subs_h2d", || { + let a = airs.to_device().unwrap(); + let s = substitutions.to_device().unwrap(); + (a, s) + }); - cuda_abi::apc_tracegen(&mut output, airs, substitutions, num_apc_calls).unwrap(); + time_stage("baseline.surviving_kernel", || { + cuda_abi::apc_tracegen(&mut output, airs, substitutions, num_apc_calls).unwrap(); + }); // Apply derived columns using the GPU expression evaluator - let (derived_specs, derived_bc) = compile_derived_to_gpu( - &self.apc.machine.derived_columns, - &apc_poly_id_to_index, - height, + let (d_specs, d_bc) = time_stage("baseline.derived_compile_h2d", || { + let (derived_specs, derived_bc) = compile_derived_to_gpu( + &self.apc.machine.derived_columns, + &apc_poly_id_to_index, + height, + ); + // In practice `d_specs` is never empty, because we will always have `is_valid` + let d_specs = derived_specs.to_device().unwrap(); + let d_bc = derived_bc.to_device().unwrap(); + (d_specs, d_bc) + }); + time_stage("baseline.derived_kernel", || { + cuda_abi::apc_apply_derived_expr(&mut output, d_specs, d_bc, num_apc_calls).unwrap(); + }); + + self.run_bus_pass(&output, num_apc_calls, &apc_poly_id_to_index, height, "baseline"); + + Some(output) + } + + /// JIT trace generation for GPU: reads DenseRecordArena bytes directly, + /// computes only surviving APC columns, bypassing full original trace generation. + fn try_generate_witness_jit( + &self, + original_arenas: OriginalArenas, + ) -> Result, OriginalArenas> { + use super::jit_mapping::{self, ArenaType, ColumnComputation}; + use crate::cuda_abi::{self, JitColumnDesc, JitInstructionDesc}; + use std::collections::BTreeMap; + + let mut original_arenas = match original_arenas { + OriginalArenas::Initialized(arenas) => arenas, + OriginalArenas::Uninitialized => { + return Ok(DeviceMatrix::dummy()); + } + }; + + let num_apc_calls = original_arenas.number_of_calls; + + // Build instruction metadata + let instructions_with_subs: Vec<_> = self + .apc + .instructions() + .zip_eq(self.apc.subs()) + .filter(|(_, subs)| !subs.is_empty()) + .collect(); + + // Get AIR names + let air_names: Vec = instructions_with_subs + .iter() + .map(|(instr, _)| self.original_airs.opcode_to_air[&instr.inner.opcode].clone()) + .collect(); + + // Build mapping tables for Dense arena type + let mappings_by_air: std::collections::HashMap<&str, jit_mapping::AirColumnMapping> = { + let mut m = std::collections::HashMap::new(); + m.insert( + "VmAirWrapper", + jit_mapping::base_alu_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::loadstore_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::shift_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::branch_equal_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::less_than_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::branch_lt_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::auipc_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::jalr_mapping_for(ArenaType::Dense), + ); + m + }; + + // Pre-check: all AIR types must have mappings + for air_name in air_names.iter() { + if !mappings_by_air.contains_key(air_name.as_str()) { + tracing::warn!("GPU JIT: no mapping for AIR '{}', falling back", air_name); + return Err(OriginalArenas::Initialized(original_arenas)); + } + } + + // Count occurrences per AIR type + let air_id_occurrences: std::collections::HashMap<&str, usize> = + air_names.iter().map(|s| s.as_str()).counts(); + + // Compute per-instruction offset within its AIR type + let instruction_offsets: Vec = air_names + .iter() + .scan( + std::collections::HashMap::<&str, usize>::default(), + |counts, air_name| { + let count = counts.entry(air_name.as_str()).or_default(); + let current = *count; + *count += 1; + Some(current) + }, + ) + .collect(); + + // Build APC poly_id to index mapping + let apc_poly_id_to_index: BTreeMap = self + .apc + .machine + .main_columns() + .enumerate() + .map(|(index, c)| (c.id, index)) + .collect(); + + let width = apc_poly_id_to_index.len(); + let height = next_power_of_two_or_zero(num_apc_calls); + + // Pre-take arenas so we can size the VPMM warm-up correctly without + // touching arena bytes inside the timed region. + let pre_taken: Vec<(String, openvm_circuit::arch::DenseRecordArena)> = air_names + .iter() + .unique() + .filter_map(|name| { + original_arenas + .take_real_arena(name) + .map(|a| (name.clone(), a)) + }) + .collect(); + let arena_total: usize = pre_taken.iter().map(|(_, a)| a.allocated().len()).sum(); + tracing::info!( + "[trace_profile] descriptor.arena_size {} bytes", + arena_total ); - // In practice `d_specs` is never empty, because we will always have `is_valid` - let d_specs = derived_specs.to_device().unwrap(); - let d_bc = derived_bc.to_device().unwrap(); - cuda_abi::apc_apply_derived_expr(&mut output, d_specs, d_bc, num_apc_calls).unwrap(); + maybe_warmup_vpmm(arena_total); - // Encode bus interactions for GPU consumption - let (bus_interactions, arg_spans, bytecode) = compile_bus_to_gpu( - &self.apc.machine.bus_interactions, - &apc_poly_id_to_index, - height, + // Direct H2D from each arena (skip host-side concatenation). + let (d_arena, arena_offsets) = time_stage("descriptor.arena_direct_h2d", || { + concat_arenas_direct_h2d(pre_taken) + }); + + // Build JIT descriptor arrays + let mut jit_instructions: Vec = Vec::new(); + let mut jit_col_descs: Vec = Vec::new(); + + // Derive range_max_bits from range checker count buffer size: + // num_rows = 1 << (range_max_bits + 1), so range_max_bits = log2(num_rows) - 1 + let range_max_bits = { + let num_bins = self.periphery.real.range_checker.count.len(); + if num_bins > 1 { (num_bins as f64).log2() as u32 - 1 } else { 29 } + }; + + time_stage("descriptor.build_descriptors", || { + for ((air_name, offset), (_, subs)) in air_names + .iter() + .zip(instruction_offsets.iter()) + .zip(instructions_with_subs.iter()) + { + let mapping = &mappings_by_air[air_name.as_str()]; + let occurrences = *air_id_occurrences.get(air_name.as_str()).unwrap(); + + let arena_offset = *arena_offsets.get(air_name).unwrap_or(&0); + let record_stride_bytes = jit_mapping::dense_record_stride(mapping.air_name); + + let col_desc_start = jit_col_descs.len() as u32; + + for sub in subs.iter() { + let apc_col_index = apc_poly_id_to_index[&sub.apc_poly_id]; + let col_mapping = &mapping.columns[sub.original_poly_index]; + jit_col_descs.push(column_comp_to_jit_desc( + &col_mapping.computation, + apc_col_index as u16, + )); + } + + let col_desc_count = jit_col_descs.len() as u32 - col_desc_start; + + jit_instructions.push(JitInstructionDesc { + arena_offset, + record_stride: (record_stride_bytes * occurrences) as u32, + record_offset: (record_stride_bytes * offset) as u32, + col_desc_start, + col_desc_count, + }); + } + }); + + // Upload descriptors to GPU + let (d_instructions, d_col_descs) = time_stage("descriptor.descriptors_h2d", || { + let d_instructions = jit_instructions.to_device().unwrap(); + let d_col_descs = jit_col_descs.to_device().unwrap(); + (d_instructions, d_col_descs) + }); + + // Allocate output (pre-zeroed) + let mut output = time_stage("descriptor.allocate_zero_output", || { + let m = DeviceMatrix::::with_capacity(height, width); + m.buffer().fill_zero().unwrap(); + m + }); + + // Launch JIT kernel (replaces Stage 1 + Stage 2) + time_stage("descriptor.surviving_kernel", || { + cuda_abi::apc_jit_tracegen( + &mut output, + &d_arena, + &d_instructions, + &d_col_descs, + num_apc_calls, + range_max_bits, + ) + .unwrap(); + }); + + // Stage 3: derived expressions + bus interactions (unchanged) + let (d_specs, d_bc) = time_stage("descriptor.derived_compile_h2d", || { + let (derived_specs, derived_bc) = compile_derived_to_gpu( + &self.apc.machine.derived_columns, + &apc_poly_id_to_index, + height, + ); + let d_specs = derived_specs.to_device().unwrap(); + let d_bc = derived_bc.to_device().unwrap(); + (d_specs, d_bc) + }); + time_stage("descriptor.derived_kernel", || { + cuda_abi::apc_apply_derived_expr(&mut output, d_specs, d_bc, num_apc_calls).unwrap(); + }); + + self.run_bus_pass(&output, num_apc_calls, &apc_poly_id_to_index, height, "descriptor"); + + Ok(output) + } + + /// NVRTC trace generation for GPU. Per-APC fallback: returns `Err` with the + /// arenas if the emitter cannot yet handle every surviving column (Phase 1 + /// supports `DirectU32` only). On success the output trace contains all + /// surviving columns plus the unchanged derived-expr / bus passes. + fn try_generate_witness_nvrtc( + &self, + original_arenas: OriginalArenas, + ) -> Result, OriginalArenas> { + use crate::powdr_extension::trace_generator::cuda::nvrtc_cache::NvrtcKernelCache; + use crate::powdr_extension::trace_generator::cuda::nvrtc_emit::{ + emit_jit_kernel_source, EmitterColumn, EmitterInput, EmitterInstruction, + }; + use crate::powdr_extension::trace_generator::jit_mapping::{ + self as jit_mapping, ArenaType, ColumnComputation, + }; + + let mut original_arenas = match original_arenas { + OriginalArenas::Initialized(arenas) => arenas, + OriginalArenas::Uninitialized => { + return Ok(DeviceMatrix::dummy()); + } + }; + + let num_apc_calls = original_arenas.number_of_calls; + + let instructions_with_subs: Vec<_> = self + .apc + .instructions() + .zip_eq(self.apc.subs()) + .filter(|(_, subs)| !subs.is_empty()) + .collect(); + + let air_names: Vec = instructions_with_subs + .iter() + .map(|(instr, _)| self.original_airs.opcode_to_air[&instr.inner.opcode].clone()) + .collect(); + + let mappings_by_air: std::collections::HashMap<&str, jit_mapping::AirColumnMapping> = { + let mut m = std::collections::HashMap::new(); + m.insert( + "VmAirWrapper", + jit_mapping::base_alu_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::loadstore_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::shift_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::branch_equal_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::less_than_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::branch_lt_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::auipc_mapping_for(ArenaType::Dense), + ); + m.insert( + "VmAirWrapper", + jit_mapping::jalr_mapping_for(ArenaType::Dense), + ); + m + }; + + // Pre-check 1: every AIR has a Dense mapping (same as descriptor path). + for air_name in air_names.iter() { + if !mappings_by_air.contains_key(air_name.as_str()) { + tracing::warn!( + "GPU NVRTC: no AIR mapping for '{}', falling back", + air_name + ); + return Err(OriginalArenas::Initialized(original_arenas)); + } + } + + let air_id_occurrences: std::collections::HashMap<&str, usize> = + air_names.iter().map(|s| s.as_str()).counts(); + + let instruction_offsets: Vec = air_names + .iter() + .scan( + std::collections::HashMap::<&str, usize>::default(), + |counts, air_name| { + let count = counts.entry(air_name.as_str()).or_default(); + let current = *count; + *count += 1; + Some(current) + }, + ) + .collect(); + + let apc_poly_id_to_index: BTreeMap = self + .apc + .machine + .main_columns() + .enumerate() + .map(|(index, c)| (c.id, index)) + .collect(); + + let width = apc_poly_id_to_index.len(); + let height = next_power_of_two_or_zero(num_apc_calls); + + // Pre-check 2 (Phase 1 emitter is DirectU32-only). Walk surviving subs + // first to fail fast before any device work; build EmitterInstruction + // list as we go. + let mut emitter_instructions: Vec = + Vec::with_capacity(instructions_with_subs.len()); + + for ((air_name, offset), (_, subs)) in air_names + .iter() + .zip(instruction_offsets.iter()) + .zip(instructions_with_subs.iter()) + { + let mapping = &mappings_by_air[air_name.as_str()]; + let occurrences = *air_id_occurrences.get(air_name.as_str()).unwrap(); + let record_stride_bytes = jit_mapping::dense_record_stride(mapping.air_name); + + let mut columns: Vec = Vec::with_capacity(subs.len()); + for sub in subs.iter() { + let col_mapping = &mapping.columns[sub.original_poly_index]; + let apc_col = apc_poly_id_to_index[&sub.apc_poly_id] as u16; + match column_comp_to_emitter(&col_mapping.computation, apc_col) { + Some(c) => columns.push(c), + None => { + tracing::debug!( + "GPU NVRTC: unsupported ColumnComputation on AIR '{}', falling back", + air_name + ); + return Err(OriginalArenas::Initialized(original_arenas)); + } + } + } + + // arena_offset is filled in below after we know the layout. + emitter_instructions.push(EmitterInstruction { + arena_offset: 0, + record_stride: (record_stride_bytes * occurrences) as u32, + record_offset: (record_stride_bytes * offset) as u32, + columns, + }); + } + + // Pre-take arenas so we can size the VPMM warm-up correctly without + // touching arena bytes inside the timed region. + let pre_taken: Vec<(String, openvm_circuit::arch::DenseRecordArena)> = air_names + .iter() + .unique() + .filter_map(|name| { + original_arenas + .take_real_arena(name) + .map(|a| (name.clone(), a)) + }) + .collect(); + let arena_total: usize = pre_taken.iter().map(|(_, a)| a.allocated().len()).sum(); + tracing::info!( + "[trace_profile] nvrtc.arena_size {} bytes", + arena_total ); - let bus_interactions = bus_interactions.to_device().unwrap(); - let arg_spans = arg_spans.to_device().unwrap(); - let bytecode = bytecode.to_device().unwrap(); + maybe_warmup_vpmm(arena_total); - // Gather GPU inputs for periphery (bus ids, count device buffers) - let periphery = &self.periphery.real; + // Direct H2D from each arena (skip host-side concatenation). + let (d_arena, arena_offsets) = time_stage("nvrtc.arena_direct_h2d", || { + concat_arenas_direct_h2d(pre_taken) + }); - // Range checker - let var_range_bus_id = self.periphery.bus_ids.range_checker as u32; - let var_range_count = &periphery.range_checker.count; + // Backfill per-instruction arena_offset. + for (i, air_name) in air_names.iter().enumerate() { + emitter_instructions[i].arena_offset = + *arena_offsets.get(air_name).unwrap_or(&0); + } - // Tuple checker - let tuple_range_checker_chip = periphery.tuple_range_checker.as_ref().unwrap(); - let tuple2_bus_id = self.periphery.bus_ids.tuple_range_checker.unwrap() as u32; - let tuple2_sizes = tuple_range_checker_chip.sizes; - let tuple2_count_u32 = tuple_range_checker_chip.count.as_ref(); + let range_max_bits = { + let num_bins = self.periphery.real.range_checker.count.len(); + if num_bins > 1 { + (num_bins as f64).log2() as u32 - 1 + } else { + 29 + } + }; - // Bitwise lookup; NUM_BITS is fixed at 8 in CUDA - let bitwise_bus_id = self.periphery.bus_ids.bitwise_lookup.unwrap() as u32; - let bitwise_count_u32 = periphery.bitwise_lookup_8.as_ref().unwrap().count.as_ref(); + let mut output = time_stage("nvrtc.allocate_zero_output", || { + let m = DeviceMatrix::::with_capacity(height, width); + m.buffer().fill_zero().unwrap(); + m + }); - // Launch GPU apply-bus to update periphery histograms on device - // Note that this is implicitly serialized after `apc_tracegen`, - // because we use the default host to device stream, which only launches - // the next kernel function after the prior (`apc_tracegen`) returns. - // This is important because bus evaluation depends on trace results. - 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(); + // Emit, compile (cached), launch. + let kernel = time_stage("nvrtc.emit_source", || { + emit_jit_kernel_source(&EmitterInput { + instructions: emitter_instructions, + }) + }); + let compiled = time_stage("nvrtc.compile_or_cache_load", || { + NvrtcKernelCache::global() + .get_or_compile(&kernel) + .expect("NVRTC compile failed") + }); - Some(output) + let block_x: u32 = 256; + let grid_x: u32 = ((num_apc_calls as u32) + block_x - 1) / block_x.max(1); + time_stage("nvrtc.surviving_kernel", || { + let rc = unsafe { + cuda_abi::powdr_nvrtc_launch_jit_v1( + compiled.function(), + output.buffer().as_mut_ptr() as *mut u32, + height, + num_apc_calls as i32, + d_arena.as_ptr(), + range_max_bits, + grid_x.max(1), + block_x, + ) + }; + assert_eq!(rc, 0, "NVRTC kernel launch failed: {}", rc); + }); + + // Optional A/B verifier: launch the descriptor kernel into a parallel + // buffer with the same arena and assert host-side equality. Fires only + // when POWDR_JIT_BACKEND_VERIFY is set; intended as a development + // guardrail while emitter coverage broadens. + if std::env::var("POWDR_JIT_BACKEND_VERIFY").is_ok() { + use crate::cuda_abi::{JitColumnDesc, JitInstructionDesc}; + use openvm_cuda_common::copy::MemCopyD2H; + + let mut jit_instructions: Vec = Vec::new(); + let mut jit_col_descs: Vec = Vec::new(); + for ((air_name, offset), (_, subs)) in air_names + .iter() + .zip(instruction_offsets.iter()) + .zip(instructions_with_subs.iter()) + { + let mapping = &mappings_by_air[air_name.as_str()]; + let occurrences = *air_id_occurrences.get(air_name.as_str()).unwrap(); + let arena_offset_inst = *arena_offsets.get(air_name).unwrap_or(&0); + let record_stride_bytes = jit_mapping::dense_record_stride(mapping.air_name); + let col_desc_start = jit_col_descs.len() as u32; + for sub in subs.iter() { + let apc_col_index = apc_poly_id_to_index[&sub.apc_poly_id]; + let col_mapping = &mapping.columns[sub.original_poly_index]; + jit_col_descs.push(column_comp_to_jit_desc( + &col_mapping.computation, + apc_col_index as u16, + )); + } + let col_desc_count = jit_col_descs.len() as u32 - col_desc_start; + jit_instructions.push(JitInstructionDesc { + arena_offset: arena_offset_inst, + record_stride: (record_stride_bytes * occurrences) as u32, + record_offset: (record_stride_bytes * offset) as u32, + col_desc_start, + col_desc_count, + }); + } + + let d_instructions = jit_instructions.to_device().unwrap(); + let d_col_descs = jit_col_descs.to_device().unwrap(); + + let mut output_desc = DeviceMatrix::::with_capacity(height, width); + output_desc.buffer().fill_zero().unwrap(); + cuda_abi::apc_jit_tracegen( + &mut output_desc, + &d_arena, + &d_instructions, + &d_col_descs, + num_apc_calls, + range_max_bits, + ) + .unwrap(); + + let host_nvrtc: Vec = output.buffer().to_host().unwrap(); + let host_desc: Vec = output_desc.buffer().to_host().unwrap(); + assert_eq!(host_nvrtc.len(), host_desc.len()); + + let mut mismatches = 0usize; + let mut first: Option<(usize, usize, u32, u32)> = None; + for col in 0..width { + for r in 0..num_apc_calls { + let off = col * height + r; + let n = host_nvrtc[off].as_canonical_u32(); + let d = host_desc[off].as_canonical_u32(); + if n != d { + mismatches += 1; + if first.is_none() { + first = Some((col, r, n, d)); + } + } + } + } + if mismatches > 0 { + let (c, r, n, d) = first.unwrap(); + panic!( + "NVRTC verifier mismatch ({}): {} cells differ. First at apc_col={} row={}: nvrtc={} descriptor={}", + kernel.name, mismatches, c, r, n, d + ); + } + tracing::info!( + "NVRTC verifier PASS: {} cols x {} rows match descriptor backend ({})", + width, num_apc_calls, kernel.name + ); + } + + // Stage 3 (unchanged): derived expressions + bus interactions. + let (d_specs, d_bc) = time_stage("nvrtc.derived_compile_h2d", || { + let (derived_specs, derived_bc) = compile_derived_to_gpu( + &self.apc.machine.derived_columns, + &apc_poly_id_to_index, + height, + ); + let d_specs = derived_specs.to_device().unwrap(); + let d_bc = derived_bc.to_device().unwrap(); + (d_specs, d_bc) + }); + time_stage("nvrtc.derived_kernel", || { + cuda_abi::apc_apply_derived_expr(&mut output, d_specs, d_bc, num_apc_calls).unwrap(); + }); + + self.run_bus_pass(&output, num_apc_calls, &apc_poly_id_to_index, height, "nvrtc"); + + Ok(output) + } +} + +/// Convert a Rust ColumnComputation into the NVRTC emitter's column variant. +/// Returns `None` for arms the emitter does not yet support, in which case +/// the caller falls back to the descriptor backend. +fn column_comp_to_emitter( + comp: &super::jit_mapping::ColumnComputation, + apc_col: u16, +) -> Option { + use crate::powdr_extension::trace_generator::cuda::nvrtc_emit::EmitterColumn; + use super::jit_mapping::ColumnComputation as CC; + match comp { + CC::DirectU32 { record_byte_offset } => Some(EmitterColumn::DirectU32 { + apc_col, + off: *record_byte_offset as u16, + }), + CC::DirectU8 { record_byte_offset } => Some(EmitterColumn::DirectU8 { + apc_col, + off: *record_byte_offset as u16, + }), + CC::DirectU16 { record_byte_offset } => Some(EmitterColumn::DirectU16 { + apc_col, + off: *record_byte_offset as u16, + }), + CC::Constant(v) => Some(EmitterColumn::Constant { + apc_col, + value: *v, + }), + // Without opcode folding we'd need to read the opcode byte at runtime; + // the host doesn't yet thread per-instruction opcode through, so + // BoolFromOpcode without folding is unsupported. + CC::BoolFromOpcode { .. } => None, + CC::Conditional { condition_byte_offset, then_comp } => { + let inner = column_comp_to_emitter(then_comp, apc_col)?; + Some(EmitterColumn::Conditional { + cond_off: *condition_byte_offset as u16, + inner: Box::new(inner), + }) + } + CC::TimestampDecomp { + curr_ts_byte_offset, + curr_ts_delta, + prev_ts_byte_offset, + limb_index, + } => Some(EmitterColumn::TimestampDecomp { + apc_col, + curr_off: *curr_ts_byte_offset as u16, + prev_off: *prev_ts_byte_offset as u16, + delta: *curr_ts_delta as u16, + limb_index: *limb_index as u16, + }), + CC::PointerLimb { + val_byte_offset, + imm_byte_offset, + imm_sign_byte_offset, + limb_index, + } => Some(EmitterColumn::PointerLimb { + apc_col, + val_off: *val_byte_offset as u16, + imm_off: *imm_byte_offset as u16, + imm_sign_off: *imm_sign_byte_offset as u16, + limb_index: *limb_index as u16, + }), + CC::AluResult { + opcode_byte_offset, + b_byte_offset, + c_byte_offset, + limb_index, + } => Some(EmitterColumn::AluResult { + apc_col, + opcode_off: *opcode_byte_offset as u16, + b_off: *b_byte_offset as u16, + c_off: *c_byte_offset as u16, + limb_index: *limb_index as u16, + }), + CC::ShiftResult { + opcode_byte_offset, + b_byte_offset, + c_byte_offset, + limb_index, + } => Some(EmitterColumn::ShiftResult { + apc_col, + opcode_off: *opcode_byte_offset as u16, + b_off: *b_byte_offset as u16, + c_off: *c_byte_offset as u16, + limb_index: *limb_index as u16, + }), + CC::ShiftBitMulLeft { opcode_byte_offset, c_byte_offset } => { + Some(EmitterColumn::ShiftBitMulLeft { + apc_col, + opcode_off: *opcode_byte_offset as u16, + c_off: *c_byte_offset as u16, + }) + } + CC::ShiftBitMulRight { opcode_byte_offset, c_byte_offset } => { + Some(EmitterColumn::ShiftBitMulRight { + apc_col, + opcode_off: *opcode_byte_offset as u16, + c_off: *c_byte_offset as u16, + }) + } + CC::ShiftBSign { opcode_byte_offset, b_byte_offset } => { + Some(EmitterColumn::ShiftBSign { + apc_col, + opcode_off: *opcode_byte_offset as u16, + b_off: *b_byte_offset as u16, + }) + } + CC::ShiftBitMarker { c_byte_offset, marker_index } => { + Some(EmitterColumn::ShiftBitMarker { + apc_col, + c_off: *c_byte_offset as u16, + marker_index: *marker_index as u16, + }) + } + CC::ShiftLimbMarker { c_byte_offset, marker_index } => { + Some(EmitterColumn::ShiftLimbMarker { + apc_col, + c_off: *c_byte_offset as u16, + marker_index: *marker_index as u16, + }) + } + CC::ShiftBitCarry { opcode_byte_offset, b_byte_offset, c_byte_offset, limb_index } => { + Some(EmitterColumn::ShiftBitCarry { + apc_col, + opcode_off: *opcode_byte_offset as u16, + b_off: *b_byte_offset as u16, + c_off: *c_byte_offset as u16, + limb_index: *limb_index as u16, + }) + } + CC::BranchEqualCmpResult { a_byte_offset, b_byte_offset, opcode_byte_offset } => { + Some(EmitterColumn::BranchEqualCmpResult { + apc_col, + a_off: *a_byte_offset as u16, + b_off: *b_byte_offset as u16, + opcode_off: *opcode_byte_offset as u16, + }) + } + CC::BranchEqualDiffInvMarker { a_byte_offset, b_byte_offset, opcode_byte_offset: _, marker_index } => { + Some(EmitterColumn::BranchEqualDiffInvMarker { + apc_col, + a_off: *a_byte_offset as u16, + b_off: *b_byte_offset as u16, + marker_index: *marker_index as u16, + }) + } + CC::LoadStoreRdRs2Ptr { rd_rs2_ptr_byte_offset } => { + Some(EmitterColumn::LoadStoreRdRs2Ptr { + apc_col, + off: *rd_rs2_ptr_byte_offset as u16, + }) + } + CC::LoadStoreNeedsWrite { rd_rs2_ptr_byte_offset } => { + Some(EmitterColumn::LoadStoreNeedsWrite { + apc_col, + off: *rd_rs2_ptr_byte_offset as u16, + }) + } + CC::LoadStoreWriteAuxPrevTs { write_prev_ts_byte_offset, rd_rs2_ptr_byte_offset } => { + Some(EmitterColumn::LoadStoreWriteAuxPrevTs { + apc_col, + write_prev_ts_off: *write_prev_ts_byte_offset as u16, + rd_rs2_ptr_off: *rd_rs2_ptr_byte_offset as u16, + }) + } + CC::LoadStoreWriteAuxDecomp { from_ts_byte_offset, write_prev_ts_byte_offset, rd_rs2_ptr_byte_offset, limb_index } => { + Some(EmitterColumn::LoadStoreWriteAuxDecomp { + apc_col, + from_ts_off: *from_ts_byte_offset as u16, + write_prev_ts_off: *write_prev_ts_byte_offset as u16, + rd_rs2_ptr_off: *rd_rs2_ptr_byte_offset as u16, + limb_index: *limb_index as u16, + }) + } + CC::LoadStoreIsLoad { opcode_byte_offset } => { + Some(EmitterColumn::LoadStoreIsLoad { + apc_col, + opcode_off: *opcode_byte_offset as u16, + }) + } + CC::LoadStoreFlag { opcode_byte_offset, shift_byte_offset, flag_index } => { + Some(EmitterColumn::LoadStoreFlag { + apc_col, + opcode_off: *opcode_byte_offset as u16, + shift_off: *shift_byte_offset as u16, + flag_index: *flag_index as u16, + }) + } + CC::LoadStoreWriteData { opcode_byte_offset, shift_byte_offset, read_data_byte_offset, prev_data_byte_offset, limb_index } => { + Some(EmitterColumn::LoadStoreWriteData { + apc_col, + opcode_off: *opcode_byte_offset as u16, + shift_off: *shift_byte_offset as u16, + read_data_off: *read_data_byte_offset as u16, + prev_data_off: *prev_data_byte_offset as u16, + limb_index: *limb_index as u16, + }) + } + // LessThan / BranchLessThan arms are not yet emitted by NVRTC. Returning + // None forces per-APC fallback to the descriptor backend, which + // implements them. + CC::LessThanCmpResult { .. } + | CC::LessThanDiffVal { .. } + | CC::LessThanDiffMarker { .. } + | CC::LessThanBMsbF { .. } + | CC::LessThanCMsbF { .. } + | CC::BranchLtCmpResult { .. } + | CC::BranchLtCmpLt { .. } + | CC::BranchLtDiffVal { .. } + | CC::BranchLtDiffMarker { .. } + | CC::BranchLtAMsbF { .. } + | CC::BranchLtBMsbF { .. } + | CC::AuipcRdLimb { .. } + | CC::JalrToPcLsb { .. } + | CC::JalrToPcLimb { .. } + | CC::JalrRdLimb { .. } + | CC::ConditionalNotMaxU32 { .. } => None, + } +} + +/// Convert a Rust ColumnComputation to a CUDA JitColumnDesc. +fn column_comp_to_jit_desc( + comp: &super::jit_mapping::ColumnComputation, + apc_col: u16, +) -> crate::cuda_abi::JitColumnDesc { + use super::jit_mapping::ColumnComputation as CC; + use crate::cuda_abi::JitColumnDesc; + + let mut desc = JitColumnDesc { + comp_type: 0, + apc_col, + p: [0; 6], + }; + + match comp { + CC::DirectU32 { record_byte_offset } => { + desc.comp_type = 0; // JIT_DIRECT_U32 + desc.p[0] = *record_byte_offset as u16; + } + CC::DirectU8 { record_byte_offset } => { + desc.comp_type = 1; // JIT_DIRECT_U8 + desc.p[0] = *record_byte_offset as u16; + } + CC::DirectU16 { record_byte_offset } => { + desc.comp_type = 2; // JIT_DIRECT_U16 + desc.p[0] = *record_byte_offset as u16; + } + CC::TimestampDecomp { curr_ts_byte_offset, curr_ts_delta, prev_ts_byte_offset, limb_index } => { + desc.comp_type = 3; // JIT_TIMESTAMP_DECOMP + desc.p[0] = *curr_ts_byte_offset as u16; + desc.p[1] = *curr_ts_delta as u16; + desc.p[2] = *prev_ts_byte_offset as u16; + desc.p[3] = *limb_index as u16; + } + CC::AluResult { opcode_byte_offset, b_byte_offset, c_byte_offset, limb_index } => { + desc.comp_type = 4; // JIT_ALU_RESULT + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *c_byte_offset as u16; + desc.p[3] = *limb_index as u16; + } + CC::BoolFromOpcode { opcode_byte_offset, expected_opcode } => { + desc.comp_type = 5; // JIT_BOOL_FROM_OPCODE + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *expected_opcode as u16; + } + CC::PointerLimb { val_byte_offset, imm_byte_offset, imm_sign_byte_offset, limb_index } => { + desc.comp_type = 6; // JIT_POINTER_LIMB + desc.p[0] = *val_byte_offset as u16; + desc.p[1] = *imm_byte_offset as u16; + desc.p[2] = *imm_sign_byte_offset as u16; + desc.p[3] = *limb_index as u16; + } + CC::Conditional { condition_byte_offset, then_comp } => { + // Encode inner computation with conditional flag + let mut inner = column_comp_to_jit_desc(then_comp, apc_col); + inner.comp_type |= 0x80; // JIT_COND_FLAG + inner.p[5] = *condition_byte_offset as u16; + return inner; + } + CC::Constant(val) => { + desc.comp_type = 16; // JIT_CONSTANT + desc.p[0] = (*val & 0xFFFF) as u16; + desc.p[1] = ((*val >> 16) & 0xFFFF) as u16; + } + CC::ShiftResult { opcode_byte_offset, b_byte_offset, c_byte_offset, limb_index } => { + desc.comp_type = 7; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *c_byte_offset as u16; + desc.p[3] = *limb_index as u16; + } + CC::ShiftBitMulLeft { opcode_byte_offset, c_byte_offset } => { + desc.comp_type = 8; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *c_byte_offset as u16; + } + CC::ShiftBitMulRight { opcode_byte_offset, c_byte_offset } => { + desc.comp_type = 9; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *c_byte_offset as u16; + } + CC::ShiftBSign { opcode_byte_offset, b_byte_offset } => { + desc.comp_type = 10; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + } + CC::ShiftBitMarker { c_byte_offset, marker_index } => { + desc.comp_type = 11; + desc.p[0] = *c_byte_offset as u16; + desc.p[1] = *marker_index as u16; + } + CC::ShiftLimbMarker { c_byte_offset, marker_index } => { + desc.comp_type = 12; + desc.p[0] = *c_byte_offset as u16; + desc.p[1] = *marker_index as u16; + } + CC::ShiftBitCarry { opcode_byte_offset, b_byte_offset, c_byte_offset, limb_index } => { + desc.comp_type = 13; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *c_byte_offset as u16; + desc.p[3] = *limb_index as u16; + } + CC::BranchEqualCmpResult { a_byte_offset, b_byte_offset, opcode_byte_offset } => { + desc.comp_type = 14; + desc.p[0] = *a_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *opcode_byte_offset as u16; + } + CC::BranchEqualDiffInvMarker { a_byte_offset, b_byte_offset, opcode_byte_offset, marker_index } => { + desc.comp_type = 15; + desc.p[0] = *a_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *opcode_byte_offset as u16; + desc.p[3] = *marker_index as u16; + } + CC::LoadStoreRdRs2Ptr { rd_rs2_ptr_byte_offset } => { + desc.comp_type = 17; + desc.p[0] = *rd_rs2_ptr_byte_offset as u16; + } + CC::LoadStoreNeedsWrite { rd_rs2_ptr_byte_offset } => { + desc.comp_type = 18; + desc.p[0] = *rd_rs2_ptr_byte_offset as u16; + } + CC::LoadStoreWriteAuxPrevTs { write_prev_ts_byte_offset, rd_rs2_ptr_byte_offset } => { + desc.comp_type = 19; + desc.p[0] = *write_prev_ts_byte_offset as u16; + desc.p[1] = *rd_rs2_ptr_byte_offset as u16; + } + CC::LoadStoreWriteAuxDecomp { from_ts_byte_offset, write_prev_ts_byte_offset, rd_rs2_ptr_byte_offset, limb_index } => { + desc.comp_type = 20; + desc.p[0] = *from_ts_byte_offset as u16; + desc.p[1] = *write_prev_ts_byte_offset as u16; + desc.p[2] = *rd_rs2_ptr_byte_offset as u16; + desc.p[3] = *limb_index as u16; + } + CC::LoadStoreIsLoad { opcode_byte_offset } => { + desc.comp_type = 21; + desc.p[0] = *opcode_byte_offset as u16; + } + CC::LoadStoreFlag { opcode_byte_offset, shift_byte_offset, flag_index } => { + desc.comp_type = 22; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *shift_byte_offset as u16; + desc.p[2] = *flag_index as u16; + } + CC::LoadStoreWriteData { opcode_byte_offset, shift_byte_offset, read_data_byte_offset, prev_data_byte_offset, limb_index } => { + desc.comp_type = 23; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *shift_byte_offset as u16; + desc.p[2] = *read_data_byte_offset as u16; + desc.p[3] = *prev_data_byte_offset as u16; + desc.p[4] = *limb_index as u16; + } + // ── LessThan arms (24-28) — share (opcode, b, c) signature ── + CC::LessThanCmpResult { opcode_byte_offset, b_byte_offset, c_byte_offset } => { + desc.comp_type = 24; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *c_byte_offset as u16; + } + CC::LessThanDiffVal { opcode_byte_offset, b_byte_offset, c_byte_offset } => { + desc.comp_type = 25; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *c_byte_offset as u16; + } + CC::LessThanDiffMarker { opcode_byte_offset, b_byte_offset, c_byte_offset, marker_index } => { + desc.comp_type = 26; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + desc.p[2] = *c_byte_offset as u16; + desc.p[3] = *marker_index as u16; + } + CC::LessThanBMsbF { opcode_byte_offset, b_byte_offset } => { + desc.comp_type = 27; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + } + CC::LessThanCMsbF { opcode_byte_offset, c_byte_offset } => { + desc.comp_type = 28; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *c_byte_offset as u16; + } + // ── BranchLessThan arms (29-34) — share (opcode, a, b) signature ── + CC::BranchLtCmpResult { opcode_byte_offset, a_byte_offset, b_byte_offset } => { + desc.comp_type = 29; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *a_byte_offset as u16; + desc.p[2] = *b_byte_offset as u16; + } + CC::BranchLtCmpLt { opcode_byte_offset, a_byte_offset, b_byte_offset } => { + desc.comp_type = 30; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *a_byte_offset as u16; + desc.p[2] = *b_byte_offset as u16; + } + CC::BranchLtDiffVal { opcode_byte_offset, a_byte_offset, b_byte_offset } => { + desc.comp_type = 31; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *a_byte_offset as u16; + desc.p[2] = *b_byte_offset as u16; + } + CC::BranchLtDiffMarker { opcode_byte_offset, a_byte_offset, b_byte_offset, marker_index } => { + desc.comp_type = 32; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *a_byte_offset as u16; + desc.p[2] = *b_byte_offset as u16; + desc.p[3] = *marker_index as u16; + } + CC::BranchLtAMsbF { opcode_byte_offset, a_byte_offset } => { + desc.comp_type = 33; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *a_byte_offset as u16; + } + CC::BranchLtBMsbF { opcode_byte_offset, b_byte_offset } => { + desc.comp_type = 34; + desc.p[0] = *opcode_byte_offset as u16; + desc.p[1] = *b_byte_offset as u16; + } + // ── Auipc ── + CC::AuipcRdLimb { pc_byte_offset, imm_byte_offset, limb_index } => { + desc.comp_type = 35; + desc.p[0] = *pc_byte_offset as u16; + desc.p[1] = *imm_byte_offset as u16; + desc.p[2] = *limb_index as u16; + } + // ── Jalr ── + CC::JalrToPcLsb { rs1_byte_offset, imm_byte_offset, imm_sign_byte_offset } => { + desc.comp_type = 36; + desc.p[0] = *rs1_byte_offset as u16; + desc.p[1] = *imm_byte_offset as u16; + desc.p[2] = *imm_sign_byte_offset as u16; + } + CC::JalrToPcLimb { rs1_byte_offset, imm_byte_offset, imm_sign_byte_offset, limb_index } => { + desc.comp_type = 37; + desc.p[0] = *rs1_byte_offset as u16; + desc.p[1] = *imm_byte_offset as u16; + desc.p[2] = *imm_sign_byte_offset as u16; + desc.p[3] = *limb_index as u16; + } + CC::JalrRdLimb { pc_byte_offset, limb_index } => { + desc.comp_type = 38; + desc.p[0] = *pc_byte_offset as u16; + desc.p[1] = *limb_index as u16; + } + CC::ConditionalNotMaxU32 { ptr_byte_offset, then_comp } => { + // Encode like Conditional but with a different flag bit (0x40 = + // JIT_COND_NOT_MAX_U32). The kernel reads p[5] as the byte offset + // of a u32 value and compares against u32::MAX. + let mut inner = column_comp_to_jit_desc(then_comp, apc_col); + inner.comp_type |= 0x40; + inner.p[5] = *ptr_byte_offset as u16; + return inner; + } } + + desc } impl>, ISA: OpenVmISA> Chip @@ -400,10 +2547,66 @@ impl>, ISA: OpenVmISA> Chip fn generate_proving_ctx(&self, _: R) -> AirProvingContext { tracing::trace!("Generating air proof input for PowdrChip {}", self.name); - let trace = self - .trace_generator - .try_generate_witness(self.record_arena_by_air_name.take()); + let backend = pick_jit_backend(); + let arenas = self.record_arena_by_air_name.take(); - AirProvingContext::new(vec![], trace, vec![]) + let trace = match backend { + JitBackend::Off => self + .trace_generator + .try_generate_witness(arenas) + .unwrap_or_else(DeviceMatrix::dummy), + JitBackend::Descriptor => match self.trace_generator.try_generate_witness_jit(arenas) { + Ok(trace) => { + tracing::info!( + "GPU JIT (descriptor) trace gen used for PowdrChip {}", + self.name + ); + trace + } + Err(arenas) => { + tracing::warn!( + "GPU JIT (descriptor) not available for PowdrChip {}, falling back", + self.name + ); + self.trace_generator + .try_generate_witness(arenas) + .unwrap_or_else(DeviceMatrix::dummy) + } + }, + JitBackend::Nvrtc => match self.trace_generator.try_generate_witness_nvrtc(arenas) { + Ok(trace) => { + tracing::info!( + "GPU JIT (nvrtc) trace gen used for PowdrChip {}", + self.name + ); + trace + } + Err(arenas) => { + tracing::debug!( + "GPU JIT (nvrtc) not yet supported for PowdrChip {}, trying descriptor", + self.name + ); + match self.trace_generator.try_generate_witness_jit(arenas) { + Ok(trace) => { + tracing::info!( + "GPU JIT (descriptor) trace gen used for PowdrChip {}", + self.name + ); + trace + } + Err(arenas) => self + .trace_generator + .try_generate_witness(arenas) + .unwrap_or_else(DeviceMatrix::dummy), + } + } + }, + }; + + AirProvingContext { + cached_mains: vec![], + common_main: trace, + public_values: vec![], + } } } diff --git a/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_bus_emit.rs b/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_bus_emit.rs new file mode 100644 index 0000000000..d125bb720b --- /dev/null +++ b/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_bus_emit.rs @@ -0,0 +1,1982 @@ +//! NVRTC bus emitter — kind-templated kernels with `__constant__` op tables. +//! +//! ## Design +//! +//! A keccak APC has ~1700 bus interactions. The earlier per-interaction +//! switch-arm design produced 30K-line kernels that NVRTC could not compile +//! within 10 minutes. This emitter takes a different approach: +//! +//! 1. **Host classifies** each bus interaction into a "simple" form that +//! fits a small fixed-size struct (`VarRangeOp`, `Tuple2Op`, +//! `BitwiseOp`). Anything that cannot be reduced to that form errors out. +//! +//! 2. **Source is fixed per kind** — never includes table values. Each kind +//! has a single `__global__` kernel template (`VAR_RANGE_KERNEL_SRC`, +//! `TUPLE2_KERNEL_SRC`, `BITWISE_RANGE_KERNEL_SRC`, +//! `BITWISE_XOR_KERNEL_SRC`). Tables flow through `__constant__` memory +//! uploaded per-APC at launch time via `cuMemcpyHtoD` to a module symbol. +//! +//! 3. **One PTX per kind**, shared across all APCs and all runs. Disk cache +//! has 4 PTX files total. Cold first compile is one-time-per-kind, then +//! every subsequent APC's bus pass is cache-warm. +//! +//! Phase 0 measurement (RTX 5090, APC=30 keccak): bytecode VM is 282 ms / +//! 90 calls = 3.13 ms/call; spike floor (no eval) is 0.25 ms/call. This +//! design targets ~0.5 ms/call by replacing the bytecode VM dispatch with +//! straight-line code that reads constant-cached op metadata. + +use std::collections::BTreeMap; + +use openvm_stark_backend::p3_field::PrimeField32; +use openvm_stark_sdk::p3_baby_bear::BabyBear; +use powdr_autoprecompiles::{ + expression::AlgebraicExpression, symbolic_machine::SymbolicBusInteraction, +}; +use powdr_expression::AlgebraicUnaryOperator; + +use super::nvrtc_emit::EmittedKernel; + +/// Bumped whenever any kernel source string changes — forces NVRTC and disk +/// PTX cache invalidation. +const EMITTER_VERSION: u32 = 7; + +/// Maximum number of column terms in a single affine arg expression. The +/// keccak APC peak observed is 5 (timestamp-delta with 5 columns); 22/23 +/// of unhandled interactions fit in 3 terms. +pub const MAX_TERMS_PER_ARG: usize = 5; + +/// Soft cap for diagnostic / safety. Op tables now live in global memory, +/// so this limit only prevents silly inputs — keccak APCs peak around 700. +const MAX_OPS_PER_KIND: usize = 16384; + +/// BabyBear prime. +const P: u32 = 0x7800_0001; + +/// Field-modular subtract for canonical inputs. `(a - b) mod P`. +fn p_sub(a: u32, b: u32) -> u32 { + if a >= b { + a - b + } else { + a + P - b + } +} + +/// Field-modular negation. `-a mod P`. +fn p_neg(a: u32) -> u32 { + if a == 0 { + 0 + } else { + P - a + } +} + +/// Sentinel for "no guard column" in the op tables — no row should ever +/// land at column `u32::MAX`, so the kernel uses this to take the +/// constant-multiplicity path. +pub const NO_GUARD: u32 = u32::MAX; + +/// One affine arg expression: `coef_const + sum_{i Self { + let mut cols = [0u32; MAX_TERMS_PER_ARG]; + let mut coefs_monty = [0u32; MAX_TERMS_PER_ARG]; + cols[0] = col; + coefs_monty[0] = host_to_monty(1); + Self { + n_terms: 1, + cols, + coefs_monty, + coef_const_monty: 0, // monty(0) = 0 + } + } +} + +/// Mult shapes (mult_const + guard_col): +/// - `Number(c)` → mult_const = c, guard_col = NO_GUARD +/// - `Reference(col)` → mult_const = 1, guard_col = col (0/1 guard) +/// - `Number(c) * Reference(col)` (any side) → mult_const = c, guard_col = col +#[repr(C)] +#[derive(Clone, Copy, Debug, Hash)] +pub struct VarRangeOp { + pub mult_const: u32, + pub guard_col: u32, + pub value: AffineArg, + pub max_bits: u32, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Hash)] +pub struct Tuple2Op { + pub mult_const: u32, + pub guard_col: u32, + pub v0: AffineArg, + pub v1: AffineArg, +} + +/// One simple-form bitwise interaction. Selector is folded by the host — +/// range and xor get separate tables and separate kernels. +#[repr(C)] +#[derive(Clone, Copy, Debug, Hash)] +pub struct BitwiseOp { + pub mult_const: u32, + pub guard_col: u32, + pub x: AffineArg, + pub y: AffineArg, +} + +/// Kind of bus interaction. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum BusKind { + VarRange, + Tuple2, + Bitwise, + Unsupported, +} + +/// One interaction's host-known shape, passed to `partition_apc_bus`. +#[derive(Clone, Debug)] +pub struct BusInteractionDesc { + pub kind: BusKind, + pub mult: AlgebraicExpression, + pub args: Vec>, +} + +#[derive(Clone, Debug)] +pub struct BusEmitterInput { + pub interactions: Vec, + pub apc_height: usize, + pub apc_poly_id_to_index: BTreeMap, +} + +/// Result of partitioning an APC's bus interactions into kind-specific +/// op tables. The affine-shape ops (`*_ops`) are used by both the +/// interpreter and codegen paths. The bilinear-shape ops (`*_ops_bilinear`) +/// are populated only when affine decode fails and bilinear succeeds; they +/// are emitted by the codegen path but treated as `unhandled` by the +/// interpreter path (which doesn't support bilinear in-kernel). +#[derive(Debug, Default)] +pub struct PartitionedBus { + pub var_ops: Vec, + pub tuple_ops: Vec, + pub bitwise_range_ops: Vec, + pub bitwise_xor_ops: Vec, + /// Bilinear-shape ops (codegen path only). Each `*_bilinear` entry + /// also carries an `unhandled` index so the interpreter path can + /// treat it as a fallback target without reparsing. + pub var_ops_bilinear: Vec<(usize, VarRangeOpBilinear)>, + pub tuple_ops_bilinear: Vec<(usize, Tuple2OpBilinear)>, + pub bitwise_range_ops_bilinear: Vec<(usize, BitwiseOpBilinear)>, + pub bitwise_xor_ops_bilinear: Vec<(usize, BitwiseOpBilinear)>, + /// Number of interactions on unsupported buses (memory/exec/program). + /// These do nothing in the existing bytecode-VM kernel either — + /// `apc_apply_bus_kernel` only updates 3 histograms. + pub n_unsupported: usize, + /// Indices into the original `bus_interactions` slice for interactions + /// whose mult/args fit neither affine nor bilinear (or whose + /// `mult` shape isn't a simple guard). The caller should run the + /// bytecode VM kernel filtered to these indices. + pub unhandled: Vec, +} + +#[derive(Debug)] +pub enum PartitionError { + /// An interaction on a supported bus has args/mult that don't reduce to + /// the simple form. The string identifies which. + NotSimple(String), + /// Wrong arity for the bus kind (e.g., var_range with 3 args). + BadArity { kind: BusKind, got: usize }, + /// Bitwise selector wasn't a constant 0/1. + BitwiseSelector(String), + /// Op-count exceeded `MAX_OPS_PER_KIND` for one kind. + TooManyOps { + kind: BusKind, + got: usize, + max: usize, + }, +} + +impl std::fmt::Display for PartitionError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PartitionError::NotSimple(s) => write!(f, "non-simple bus interaction: {}", s), + PartitionError::BadArity { kind, got } => { + write!(f, "bad arity for {:?}: got {} args", kind, got) + } + PartitionError::BitwiseSelector(s) => { + write!(f, "non-constant bitwise selector: {}", s) + } + PartitionError::TooManyOps { kind, got, max } => { + write!(f, "too many {:?} ops: {} > MAX={}", kind, got, max) + } + } + } +} + +impl std::error::Error for PartitionError {} + +/// Classify a bus interaction id against the known periphery bus ids. +pub fn classify( + id: u64, + var_range_bus_id: u64, + tuple2_bus_id: Option, + bitwise_bus_id: Option, +) -> BusKind { + if id == var_range_bus_id { + BusKind::VarRange + } else if Some(id) == tuple2_bus_id { + BusKind::Tuple2 + } else if Some(id) == bitwise_bus_id { + BusKind::Bitwise + } else { + BusKind::Unsupported + } +} + +/// Build a `BusEmitterInput` from the powdr machine's bus interactions and +/// the periphery bus ids. +pub fn build_emitter_input( + bus_interactions: &[SymbolicBusInteraction], + apc_poly_id_to_index: &BTreeMap, + apc_height: usize, + var_range_bus_id: u64, + tuple2_bus_id: Option, + bitwise_bus_id: Option, +) -> BusEmitterInput { + let interactions = bus_interactions + .iter() + .map(|bi| BusInteractionDesc { + kind: classify(bi.id as u64, var_range_bus_id, tuple2_bus_id, bitwise_bus_id), + mult: bi.mult.clone(), + args: bi.args.clone(), + }) + .collect(); + + BusEmitterInput { + interactions, + apc_height, + apc_poly_id_to_index: apc_poly_id_to_index.clone(), + } +} + +/// Helper: extract a u32 if the expression is `Number(c)`. +fn as_const_u32(expr: &AlgebraicExpression) -> Option { + if let AlgebraicExpression::Number(c) = expr { + Some(c.as_canonical_u32()) + } else { + None + } +} + +/// Helper: extract a column index if the expression is `Reference(id)`. +fn as_col( + expr: &AlgebraicExpression, + apc_poly_id_to_index: &BTreeMap, +) -> Option { + if let AlgebraicExpression::Reference(r) = expr { + apc_poly_id_to_index.get(&r.id).map(|c| *c as u32) + } else { + None + } +} + +/// Pre-simplify shallow mul-by-1, mul-by-0 / unary-Minus into canonical +/// form. Folds `c * x` and `x * c` into `c * x` (constant always on +/// left) so the partition matcher has fewer cases. Recursive only on +/// children, not deeper than needed for the patterns we accept. +fn simplify_mult( + expr: &AlgebraicExpression, +) -> std::borrow::Cow<'_, AlgebraicExpression> { + use std::borrow::Cow; + + if let AlgebraicExpression::BinaryOperation(b) = expr { + if matches!(b.op, powdr_expression::AlgebraicBinaryOperator::Mul) { + let lhs = simplify_mult(&b.left); + let rhs = simplify_mult(&b.right); + if let AlgebraicExpression::Number(c) = lhs.as_ref() { + if c.as_canonical_u32() == 1 { + return Cow::Owned((*rhs).clone()); + } + if c.as_canonical_u32() == 0 { + return Cow::Owned(AlgebraicExpression::Number(*c)); + } + } + if let AlgebraicExpression::Number(c) = rhs.as_ref() { + if c.as_canonical_u32() == 1 { + return Cow::Owned((*lhs).clone()); + } + if c.as_canonical_u32() == 0 { + return Cow::Owned(AlgebraicExpression::Number(*c)); + } + } + } + } + std::borrow::Cow::Borrowed(expr) +} + +/// Result of decoding an arg into multi-term affine form. Terms are +/// (col, coef_canon) pairs in canonical (non-Montgomery) BabyBear; the +/// constant offset is also canonical. Caller converts to Monty when +/// building the op table. +#[derive(Debug, Default)] +struct AffineDecoded { + /// Up to MAX_TERMS_PER_ARG distinct (col, coef) pairs. Same-column terms + /// are merged at decode time so we don't waste a slot. + terms: Vec<(u32, u32)>, + /// Sum of all Number leaves multiplied through the surrounding factors. + constant: u32, +} + +/// Recursive walker that flattens a nested Add/Sub/Mul/Minus expression +/// into AffineDecoded.terms + constant. `factor` is the canonical-field +/// scalar multiplier from outer context (combines Number-Mul nesting and +/// Minus negation). Returns false if the expression can't be expressed +/// affinely (e.g., contains a Reference*Reference). +fn collect_affine( + expr: &AlgebraicExpression, + apc_poly_id_to_index: &BTreeMap, + factor: u32, + out: &mut AffineDecoded, +) -> bool { + match expr { + AlgebraicExpression::Number(c) => { + let v = (factor as u64) * (c.as_canonical_u32() as u64) % (P as u64); + out.constant = ((out.constant as u64 + v) % (P as u64)) as u32; + true + } + AlgebraicExpression::Reference(r) => { + let col = match apc_poly_id_to_index.get(&r.id) { + Some(c) => *c as u32, + None => return false, + }; + // Merge with existing same-col term if present. + if let Some(existing) = out.terms.iter_mut().find(|(c, _)| *c == col) { + let new_coef = (existing.1 as u64 + factor as u64) % (P as u64); + existing.1 = new_coef as u32; + } else { + out.terms.push((col, factor)); + } + true + } + AlgebraicExpression::UnaryOperation(u) => { + if !matches!(u.op, AlgebraicUnaryOperator::Minus) { + return false; + } + collect_affine(&u.expr, apc_poly_id_to_index, p_neg(factor), out) + } + AlgebraicExpression::BinaryOperation(b) => match b.op { + powdr_expression::AlgebraicBinaryOperator::Add => { + collect_affine(&b.left, apc_poly_id_to_index, factor, out) + && collect_affine(&b.right, apc_poly_id_to_index, factor, out) + } + powdr_expression::AlgebraicBinaryOperator::Sub => { + collect_affine(&b.left, apc_poly_id_to_index, factor, out) + && collect_affine( + &b.right, + apc_poly_id_to_index, + p_neg(factor), + out, + ) + } + powdr_expression::AlgebraicBinaryOperator::Mul => { + // One side must be Number for the result to stay affine. + if let Some(c) = as_const_u32(&b.left) { + let new_factor = (factor as u64 * c as u64) % (P as u64); + return collect_affine( + &b.right, + apc_poly_id_to_index, + new_factor as u32, + out, + ); + } + if let Some(c) = as_const_u32(&b.right) { + let new_factor = (factor as u64 * c as u64) % (P as u64); + return collect_affine( + &b.left, + apc_poly_id_to_index, + new_factor as u32, + out, + ); + } + false + } + }, + } +} + +/// Decode an arg expression into multi-term affine form. Returns None if +/// the expression isn't affine (e.g., Reference*Reference) or if the term +/// count would exceed `MAX_TERMS_PER_ARG`. Always omits zero-coefficient +/// terms so the op-table is dense. +fn decode_affine_arg( + expr: &AlgebraicExpression, + apc_poly_id_to_index: &BTreeMap, +) -> Option { + let mut decoded = AffineDecoded::default(); + if !collect_affine(expr, apc_poly_id_to_index, 1, &mut decoded) { + return None; + } + // Drop any zero-coef terms (rare but possible after merging). + decoded.terms.retain(|(_, c)| *c != 0); + if decoded.terms.len() > MAX_TERMS_PER_ARG { + return None; + } + let mut cols = [0u32; MAX_TERMS_PER_ARG]; + let mut coefs_monty = [0u32; MAX_TERMS_PER_ARG]; + for (i, (col, coef)) in decoded.terms.iter().enumerate() { + cols[i] = *col; + coefs_monty[i] = host_to_monty(*coef); + } + Some(AffineArg { + n_terms: decoded.terms.len() as u32, + cols, + coefs_monty, + coef_const_monty: host_to_monty(decoded.constant), + }) +} + +// ============================================================================ +// Bilinear decoder — captures degree ≤ 2 expressions that affine can't. +// +// Represents an expression as a polynomial in column references with degree +// at most 2. Internally tracks (constant, [linear terms], [bilinear terms]) +// in canonical (non-Monty) form. Recursive walker reduces every Add/Sub/Mul +// pair to a sum of monomials and rejects if any product would yield degree +// > 2. +// +// Used ONLY by the codegen path (POWDR_BUS_CODEGEN=1). The interpreter +// path keeps the simple AffineArg shape; bilinear ops fall to its +// unhandled tail there. +// ============================================================================ + +/// Polynomial-of-references with degree ≤ 2. All coefs in canonical form. +/// Linear terms keyed by single col; bilinear terms by unordered (col_a, +/// col_b) pair (a ≤ b normalized at insertion). +#[derive(Clone, Debug, Default, Hash)] +pub struct BilinearMonomials { + pub constant: u32, + pub linear: Vec<(u32, u32)>, // (col, coef_canon) + pub bilinear: Vec<(u32, u32, u32)>, // (col_a, col_b, coef_canon), col_a ≤ col_b +} + +impl BilinearMonomials { + fn from_const(c: u32) -> Self { + Self { + constant: c, + ..Default::default() + } + } + + fn from_col(col: u32) -> Self { + let mut m = Self::default(); + m.linear.push((col, 1)); + m + } + + fn add_const(&mut self, c: u32) { + self.constant = ((self.constant as u64 + c as u64) % (P as u64)) as u32; + } + + fn add_linear(&mut self, col: u32, coef: u32) { + if coef == 0 { + return; + } + if let Some(existing) = self.linear.iter_mut().find(|(c, _)| *c == col) { + existing.1 = ((existing.1 as u64 + coef as u64) % (P as u64)) as u32; + } else { + self.linear.push((col, coef)); + } + } + + fn add_bilinear(&mut self, mut col_a: u32, mut col_b: u32, coef: u32) { + if coef == 0 { + return; + } + if col_a > col_b { + std::mem::swap(&mut col_a, &mut col_b); + } + if let Some(existing) = self + .bilinear + .iter_mut() + .find(|(a, b, _)| *a == col_a && *b == col_b) + { + existing.2 = ((existing.2 as u64 + coef as u64) % (P as u64)) as u32; + } else { + self.bilinear.push((col_a, col_b, coef)); + } + } + + fn add(self, other: Self) -> Self { + let mut out = self; + out.add_const(other.constant); + for (c, k) in other.linear { + out.add_linear(c, k); + } + for (a, b, k) in other.bilinear { + out.add_bilinear(a, b, k); + } + out + } + + fn negate(mut self) -> Self { + self.constant = p_neg(self.constant); + for (_, k) in &mut self.linear { + *k = p_neg(*k); + } + for (_, _, k) in &mut self.bilinear { + *k = p_neg(*k); + } + self + } + + fn scalar_mul(mut self, c: u32) -> Self { + if c == 0 { + return Self::default(); + } + self.constant = ((self.constant as u64 * c as u64) % (P as u64)) as u32; + for (_, k) in &mut self.linear { + *k = ((*k as u64 * c as u64) % (P as u64)) as u32; + } + for (_, _, k) in &mut self.bilinear { + *k = ((*k as u64 * c as u64) % (P as u64)) as u32; + } + self + } + + /// Multiply two polynomials. Returns None if the result would exceed + /// degree 2 (i.e., one side has bilinear terms AND the other has a + /// non-constant term). + fn mul(self, other: Self) -> Option { + // Quick paths: pure constant on either side. + if self.linear.is_empty() && self.bilinear.is_empty() { + return Some(other.scalar_mul(self.constant)); + } + if other.linear.is_empty() && other.bilinear.is_empty() { + return Some(self.scalar_mul(other.constant)); + } + // Either side has linear or bilinear. If one side has bilinear + // and the other has any non-constant, result is degree > 2. + let self_has_bilinear = !self.bilinear.is_empty(); + let other_has_bilinear = !other.bilinear.is_empty(); + let self_has_nonconst = !self.linear.is_empty() || self_has_bilinear; + let other_has_nonconst = !other.linear.is_empty() || other_has_bilinear; + if (self_has_bilinear && other_has_nonconst) + || (other_has_bilinear && self_has_nonconst) + { + return None; + } + // Safe: at most linear × linear + linear × constant + constant × constant. + let mut out = Self::default(); + // const × const + out.add_const(((self.constant as u64 * other.constant as u64) % (P as u64)) as u32); + // const × linear + for (col, k) in &other.linear { + out.add_linear( + *col, + ((self.constant as u64 * *k as u64) % (P as u64)) as u32, + ); + } + for (col, k) in &self.linear { + out.add_linear( + *col, + ((other.constant as u64 * *k as u64) % (P as u64)) as u32, + ); + } + // linear × linear → bilinear (or square → linear-of-square which we + // store as bilinear with col_a == col_b). + for (col_a, k_a) in &self.linear { + for (col_b, k_b) in &other.linear { + let coef = ((*k_a as u64 * *k_b as u64) % (P as u64)) as u32; + out.add_bilinear(*col_a, *col_b, coef); + } + } + // const × bilinear (only one side has bilinear; we already + // verified the other has no non-const, but double-check + // syntactically below). + for (a, b, k) in &self.bilinear { + out.add_bilinear( + *a, + *b, + ((other.constant as u64 * *k as u64) % (P as u64)) as u32, + ); + } + for (a, b, k) in &other.bilinear { + out.add_bilinear( + *a, + *b, + ((self.constant as u64 * *k as u64) % (P as u64)) as u32, + ); + } + Some(out) + } +} + +/// Reduce an algebraic expression to a degree-≤2 polynomial of column +/// references. Returns None if the expression has degree > 2 anywhere +/// (e.g., `col_a * col_b * col_c`). +fn to_bilinear_monomials( + expr: &AlgebraicExpression, + map: &BTreeMap, +) -> Option { + match expr { + AlgebraicExpression::Number(c) => Some(BilinearMonomials::from_const(c.as_canonical_u32())), + AlgebraicExpression::Reference(r) => { + let col = *map.get(&r.id)? as u32; + Some(BilinearMonomials::from_col(col)) + } + AlgebraicExpression::UnaryOperation(u) => { + if !matches!(u.op, AlgebraicUnaryOperator::Minus) { + return None; + } + to_bilinear_monomials(&u.expr, map).map(|m| m.negate()) + } + AlgebraicExpression::BinaryOperation(b) => { + let l = to_bilinear_monomials(&b.left, map)?; + let r = to_bilinear_monomials(&b.right, map)?; + match b.op { + powdr_expression::AlgebraicBinaryOperator::Add => Some(l.add(r)), + powdr_expression::AlgebraicBinaryOperator::Sub => Some(l.add(r.negate())), + powdr_expression::AlgebraicBinaryOperator::Mul => l.mul(r), + } + } + } +} + +/// Decode an arg expression as a bilinear (degree ≤ 2) polynomial. +/// Returns None if degree > 2. Returns a "trivially affine" decode (no +/// bilinear terms) too — caller should prefer `decode_affine_arg` first +/// for the canonical compact form. +pub fn decode_bilinear_arg( + expr: &AlgebraicExpression, + apc_poly_id_to_index: &BTreeMap, +) -> Option { + to_bilinear_monomials(expr, apc_poly_id_to_index) +} + +// ============================================================================ +// Bilinear ops (codegen path only) +// ============================================================================ + +/// var_range op with bilinear value expression. Codegen emits the value +/// computation as `c_const + sum (c_lin * col) + sum (c_bil * col_a * col_b)`. +#[derive(Clone, Debug, Hash)] +pub struct VarRangeOpBilinear { + pub mult_const: u32, + pub guard_col: u32, + pub value: BilinearMonomials, + pub max_bits: u32, +} + +#[derive(Clone, Debug, Hash)] +pub struct Tuple2OpBilinear { + pub mult_const: u32, + pub guard_col: u32, + pub v0: BilinearMonomials, + pub v1: BilinearMonomials, +} + +#[derive(Clone, Debug, Hash)] +pub struct BitwiseOpBilinear { + pub mult_const: u32, + pub guard_col: u32, + pub x: BilinearMonomials, + pub y: BilinearMonomials, +} + +/// Decode a mult expression into `(mult_const, guard_col)` pair. +/// +/// - `Number(c)` → (c, NO_GUARD): unconditional multiplicity c +/// - `Reference(col)` → (1, col): 0/1 guard, single count when on +/// - `Number(c) * Reference(col)` (after simplification) → (c, col) +/// - anything else → None (caller errors out) +fn decode_mult( + expr: &AlgebraicExpression, + apc_poly_id_to_index: &BTreeMap, +) -> Option<(u32, u32)> { + let simplified = simplify_mult(expr); + let e: &AlgebraicExpression = simplified.as_ref(); + if let Some(c) = as_const_u32(e) { + return Some((c, NO_GUARD)); + } + if let Some(col) = as_col(e, apc_poly_id_to_index) { + return Some((1, col)); + } + if let AlgebraicExpression::BinaryOperation(b) = e { + if matches!(b.op, powdr_expression::AlgebraicBinaryOperator::Mul) { + // After simplify, neither side is Number(0) or Number(1). + // Look for Number(c) * Reference(col). + if let (Some(c), Some(col)) = ( + as_const_u32(&b.left), + as_col(&b.right, apc_poly_id_to_index), + ) { + return Some((c, col)); + } + if let (Some(col), Some(c)) = ( + as_col(&b.left, apc_poly_id_to_index), + as_const_u32(&b.right), + ) { + return Some((c, col)); + } + } + } + None +} + +/// Partition the input into kind-specific op tables. Interactions whose +/// mult or args don't fit the simple form land in `unhandled` for +/// bytecode-VM fallback. Errors only on truly broken inputs (bad arity, +/// non-const bitwise selector, op-count exceeding `MAX_OPS_PER_KIND`). +pub fn partition_apc_bus(input: &BusEmitterInput) -> Result { + let mut p = PartitionedBus::default(); + + for (i, intr) in input.interactions.iter().enumerate() { + match intr.kind { + BusKind::Unsupported => { + p.n_unsupported += 1; + } + BusKind::VarRange => { + if intr.args.len() != 2 { + return Err(PartitionError::BadArity { + kind: intr.kind, + got: intr.args.len(), + }); + } + let mult = decode_mult(&intr.mult, &input.apc_poly_id_to_index); + let max_bits = as_const_u32(&intr.args[1]); + let value = decode_affine_arg(&intr.args[0], &input.apc_poly_id_to_index); + + if let (Some((mult_const, guard_col)), Some(value), Some(max_bits)) = + (mult, value, max_bits) + { + if mult_const == 0 { + continue; + } + p.var_ops.push(VarRangeOp { + mult_const, + guard_col, + value, + max_bits, + }); + continue; + } + // Try bilinear before falling to unhandled. + if let (Some((mult_const, guard_col)), Some(max_bits)) = (mult, max_bits) { + if mult_const == 0 { + continue; + } + if let Some(value_bil) = + decode_bilinear_arg(&intr.args[0], &input.apc_poly_id_to_index) + { + p.var_ops_bilinear.push(( + i, + VarRangeOpBilinear { + mult_const, + guard_col, + value: value_bil, + max_bits, + }, + )); + continue; + } + } + p.unhandled.push(i); + } + BusKind::Tuple2 => { + if intr.args.len() != 2 { + return Err(PartitionError::BadArity { + kind: intr.kind, + got: intr.args.len(), + }); + } + let mult = decode_mult(&intr.mult, &input.apc_poly_id_to_index); + let v0_aff = decode_affine_arg(&intr.args[0], &input.apc_poly_id_to_index); + let v1_aff = decode_affine_arg(&intr.args[1], &input.apc_poly_id_to_index); + if let (Some((mult_const, guard_col)), Some(v0), Some(v1)) = + (mult, v0_aff, v1_aff) + { + if mult_const == 0 { + continue; + } + p.tuple_ops.push(Tuple2Op { + mult_const, + guard_col, + v0, + v1, + }); + continue; + } + if let Some((mult_const, guard_col)) = mult { + if mult_const == 0 { + continue; + } + let v0 = + decode_bilinear_arg(&intr.args[0], &input.apc_poly_id_to_index); + let v1 = + decode_bilinear_arg(&intr.args[1], &input.apc_poly_id_to_index); + if let (Some(v0), Some(v1)) = (v0, v1) { + p.tuple_ops_bilinear.push(( + i, + Tuple2OpBilinear { + mult_const, + guard_col, + v0, + v1, + }, + )); + continue; + } + } + p.unhandled.push(i); + } + BusKind::Bitwise => { + if intr.args.len() != 4 { + return Err(PartitionError::BadArity { + kind: intr.kind, + got: intr.args.len(), + }); + } + let mult = decode_mult(&intr.mult, &input.apc_poly_id_to_index); + let x_aff = decode_affine_arg(&intr.args[0], &input.apc_poly_id_to_index); + let y_aff = decode_affine_arg(&intr.args[1], &input.apc_poly_id_to_index); + let selector = as_const_u32(&intr.args[3]).ok_or_else(|| { + PartitionError::BitwiseSelector(format!("{:?}", intr.args[3])) + })?; + if let (Some((mult_const, guard_col)), Some(x), Some(y)) = + (mult, x_aff, y_aff) + { + if mult_const == 0 { + continue; + } + let op = BitwiseOp { + mult_const, + guard_col, + x, + y, + }; + match selector { + 0 => p.bitwise_range_ops.push(op), + 1 => p.bitwise_xor_ops.push(op), + other => { + return Err(PartitionError::BitwiseSelector(format!( + "selector const = {}", + other + ))) + } + } + continue; + } + // Bilinear fallback for bitwise. + if let Some((mult_const, guard_col)) = mult { + if mult_const == 0 { + continue; + } + let x = decode_bilinear_arg(&intr.args[0], &input.apc_poly_id_to_index); + let y = decode_bilinear_arg(&intr.args[1], &input.apc_poly_id_to_index); + if let (Some(x), Some(y)) = (x, y) { + let op_bil = BitwiseOpBilinear { + mult_const, + guard_col, + x, + y, + }; + match selector { + 0 => p.bitwise_range_ops_bilinear.push((i, op_bil)), + 1 => p.bitwise_xor_ops_bilinear.push((i, op_bil)), + other => { + return Err(PartitionError::BitwiseSelector(format!( + "selector const = {}", + other + ))) + } + } + continue; + } + } + p.unhandled.push(i); + } + } + } + + // Cap each kind so callers can rely on the `__constant__` array size. + if p.var_ops.len() > MAX_OPS_PER_KIND { + return Err(PartitionError::TooManyOps { + kind: BusKind::VarRange, + got: p.var_ops.len(), + max: MAX_OPS_PER_KIND, + }); + } + if p.tuple_ops.len() > MAX_OPS_PER_KIND { + return Err(PartitionError::TooManyOps { + kind: BusKind::Tuple2, + got: p.tuple_ops.len(), + max: MAX_OPS_PER_KIND, + }); + } + let bw_total = p.bitwise_range_ops.len() + p.bitwise_xor_ops.len(); + if bw_total > MAX_OPS_PER_KIND { + return Err(PartitionError::TooManyOps { + kind: BusKind::Bitwise, + got: bw_total, + max: MAX_OPS_PER_KIND, + }); + } + + Ok(p) +} + +/// Returns the four kernel-source strings that all bus passes share, in +/// emission order: var_range, tuple2, bitwise_range, bitwise_xor. +pub fn kernel_sources() -> [(&'static str, EmittedKernel); 4] { + [ + ( + "k_var_ops", + make_kernel("apc_bus_var_range", VAR_RANGE_KERNEL_SRC, 1), + ), + ( + "k_tup_ops", + make_kernel("apc_bus_tuple2", TUPLE2_KERNEL_SRC, 2), + ), + ( + "k_bit_range_ops", + make_kernel("apc_bus_bitwise_range", BITWISE_RANGE_KERNEL_SRC, 3), + ), + ( + "k_bit_xor_ops", + make_kernel("apc_bus_bitwise_xor", BITWISE_XOR_KERNEL_SRC, 4), + ), + ] +} + +fn make_kernel(name: &str, src: &str, kind_tag: u32) -> EmittedKernel { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + + let full = format!("{}{}", BUS_PRELUDE, src); + let mut hasher = DefaultHasher::new(); + "BUS_TEMPLATED".hash(&mut hasher); + EMITTER_VERSION.hash(&mut hasher); + kind_tag.hash(&mut hasher); + full.hash(&mut hasher); + let source_hash = hasher.finish(); + + EmittedKernel { + source: full, + name: name.to_string(), + source_hash, + } +} + +// ============================================================================ +// Per-APC codegen path +// +// Each op is emitted as a `case` block with constants baked as immediates. +// ptxas can constant-fold `(1 << max_bits)`, eliminate the unused-term loop +// of the interpreter approach, hoist `col * H` out of the row loop, and +// fuse mul+add into FMA. Per-row SASS drops from ~50 (interpreter) to +// ~15-20 (codegen) for the common 1-3 term affine cases. +// +// Source size: ~10 lines per op × N ops per kind. For keccak APCs the +// peak is ~700 ops in a single kind, so ~7K lines per kernel. Kernel +// source compiles in ~1-3s cold; warm via PTX disk cache is ~1ms. +// ============================================================================ + +/// Emit a per-APC `var_range` codegen kernel. Returns `None` if no ops +/// (caller should skip the kernel entirely). +pub fn emit_codegen_var_range( + affine_ops: &[VarRangeOp], + bilinear_ops: &[(usize, VarRangeOpBilinear)], +) -> Option { + if affine_ops.is_empty() && bilinear_ops.is_empty() { + return None; + } + let n_ops = affine_ops.len() + bilinear_ops.len(); + + let mut s = String::new(); + s.push_str(BUS_PRELUDE); + + let source_hash = { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut h = DefaultHasher::new(); + "BUS_CODEGEN".hash(&mut h); + EMITTER_VERSION.hash(&mut h); + "VAR_RANGE_CODEGEN".hash(&mut h); + affine_ops.hash(&mut h); + bilinear_ops.hash(&mut h); + h.finish() + }; + let name = format!("apc_bus_var_range_codegen_{:016x}", source_hash); + + use std::fmt::Write as _; + writeln!( + s, + r#"extern "C" __global__ void {name}( + const unsigned int* __restrict__ d_output, + int N, + unsigned long long H, + unsigned int* __restrict__ d_var_hist, + unsigned int var_num_bins) +{{ + const int warp = (threadIdx.x >> 5); + const int lane = (threadIdx.x & 31); + const int wpb = (blockDim.x >> 5); + constexpr int N_OPS = {n_ops}; + lookup::Histogram hist(d_var_hist, var_num_bins); + + for (int i = blockIdx.x * wpb + warp; i < N_OPS; i += gridDim.x * wpb) {{ + switch (i) {{"#, + ) + .unwrap(); + + let mut idx = 0usize; + for op in affine_ops { + writeln!(s, " case {}: {{", idx).unwrap(); + emit_codegen_var_range_body(&mut s, op); + s.push_str(" break;\n }\n"); + idx += 1; + } + for (_orig_i, op) in bilinear_ops { + writeln!(s, " case {}: {{ /* bilinear */", idx).unwrap(); + emit_codegen_var_range_bilinear_body(&mut s, op); + s.push_str(" break;\n }\n"); + idx += 1; + } + s.push_str( + " default: break; + } + } +} +", + ); + + Some(EmittedKernel { + source: s, + name, + source_hash, + }) +} + +fn emit_codegen_var_range_bilinear_body(s: &mut String, op: &VarRangeOpBilinear) { + use std::fmt::Write as _; + let bits_one = 1u32 << op.max_bits; + s.push_str(" for (int r = lane; r < N; r += 32) {\n"); + if op.guard_col != NO_GUARD { + writeln!( + s, + " unsigned int g = monty_reduce(d_output[{}ull * H + (unsigned long long)r]); + if (g == 0u) continue;", + op.guard_col + ) + .unwrap(); + } + emit_bilinear_inline(s, &op.value, " ", "v"); + writeln!( + s, + " unsigned int idx = {bits_one}u + v - 1u;" + ) + .unwrap(); + if op.mult_const == 1 { + s.push_str(" hist.add_count(idx);\n"); + } else { + writeln!( + s, + " hist.add_count_n(idx, {}u);", + op.mult_const + ) + .unwrap(); + } + s.push_str(" }\n"); +} + +fn emit_codegen_var_range_body(s: &mut String, op: &VarRangeOp) { + use std::fmt::Write as _; + let bits_one = 1u32 << op.max_bits; + s.push_str(" for (int r = lane; r < N; r += 32) {\n"); + if op.guard_col != NO_GUARD { + writeln!( + s, + " unsigned int g = monty_reduce(d_output[{}ull * H + (unsigned long long)r]); + if (g == 0u) continue;", + op.guard_col + ) + .unwrap(); + } + emit_affine_inline(s, &op.value, " ", "v"); + writeln!( + s, + " unsigned int idx = {bits_one}u + v - 1u;" + ) + .unwrap(); + if op.mult_const == 1 { + s.push_str(" hist.add_count(idx);\n"); + } else { + writeln!( + s, + " hist.add_count_n(idx, {}u);", + op.mult_const + ) + .unwrap(); + } + s.push_str(" }\n"); +} + +pub fn emit_codegen_tuple2( + affine_ops: &[Tuple2Op], + bilinear_ops: &[(usize, Tuple2OpBilinear)], +) -> Option { + if affine_ops.is_empty() && bilinear_ops.is_empty() { + return None; + } + let n_ops = affine_ops.len() + bilinear_ops.len(); + let mut s = String::new(); + s.push_str(BUS_PRELUDE); + let source_hash = { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut h = DefaultHasher::new(); + "BUS_CODEGEN".hash(&mut h); + EMITTER_VERSION.hash(&mut h); + "TUPLE2_CODEGEN".hash(&mut h); + affine_ops.hash(&mut h); + bilinear_ops.hash(&mut h); + h.finish() + }; + let name = format!("apc_bus_tuple2_codegen_{:016x}", source_hash); + + use std::fmt::Write as _; + writeln!( + s, + r#"extern "C" __global__ void {name}( + const unsigned int* __restrict__ d_output, + int N, + unsigned long long H, + unsigned int* __restrict__ d_tuple2_hist, + unsigned int tuple2_sz0, + unsigned int tuple2_sz1) +{{ + const int warp = (threadIdx.x >> 5); + const int lane = (threadIdx.x & 31); + const int wpb = (blockDim.x >> 5); + constexpr int N_OPS = {n_ops}; + const unsigned int total_bins = tuple2_sz0 * tuple2_sz1; + lookup::Histogram hist(d_tuple2_hist, total_bins); + + for (int i = blockIdx.x * wpb + warp; i < N_OPS; i += gridDim.x * wpb) {{ + switch (i) {{"#, + ) + .unwrap(); + + let mut idx = 0usize; + for op in affine_ops { + writeln!(s, " case {}: {{", idx).unwrap(); + s.push_str(" for (int r = lane; r < N; r += 32) {\n"); + if op.guard_col != NO_GUARD { + writeln!( + s, + " unsigned int g = monty_reduce(d_output[{}ull * H + (unsigned long long)r]); + if (g == 0u) continue;", + op.guard_col + ) + .unwrap(); + } + emit_affine_inline(&mut s, &op.v0, " ", "v0"); + emit_affine_inline(&mut s, &op.v1, " ", "v1"); + s.push_str(" unsigned int idx = v0 * tuple2_sz1 + v1;\n"); + if op.mult_const == 1 { + s.push_str(" hist.add_count(idx);\n"); + } else { + writeln!( + s, + " hist.add_count_n(idx, {}u);", + op.mult_const + ) + .unwrap(); + } + s.push_str(" }\n break;\n }\n"); + idx += 1; + } + for (_orig_i, op) in bilinear_ops { + writeln!(s, " case {}: {{ /* bilinear */", idx).unwrap(); + s.push_str(" for (int r = lane; r < N; r += 32) {\n"); + if op.guard_col != NO_GUARD { + writeln!( + s, + " unsigned int g = monty_reduce(d_output[{}ull * H + (unsigned long long)r]); + if (g == 0u) continue;", + op.guard_col + ) + .unwrap(); + } + emit_bilinear_inline(&mut s, &op.v0, " ", "v0"); + emit_bilinear_inline(&mut s, &op.v1, " ", "v1"); + s.push_str(" unsigned int idx = v0 * tuple2_sz1 + v1;\n"); + if op.mult_const == 1 { + s.push_str(" hist.add_count(idx);\n"); + } else { + writeln!( + s, + " hist.add_count_n(idx, {}u);", + op.mult_const + ) + .unwrap(); + } + s.push_str(" }\n break;\n }\n"); + idx += 1; + } + s.push_str(" default: break;\n }\n }\n}\n"); + + Some(EmittedKernel { + source: s, + name, + source_hash, + }) +} + +pub fn emit_codegen_bitwise( + affine_ops: &[BitwiseOp], + bilinear_ops: &[(usize, BitwiseOpBilinear)], + is_xor: bool, +) -> Option { + if affine_ops.is_empty() && bilinear_ops.is_empty() { + return None; + } + let n_ops = affine_ops.len() + bilinear_ops.len(); + let mut s = String::new(); + s.push_str(BUS_PRELUDE); + let tag = if is_xor { + "BITWISE_XOR_CODEGEN" + } else { + "BITWISE_RANGE_CODEGEN" + }; + let source_hash = { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut h = DefaultHasher::new(); + "BUS_CODEGEN".hash(&mut h); + EMITTER_VERSION.hash(&mut h); + tag.hash(&mut h); + affine_ops.hash(&mut h); + bilinear_ops.hash(&mut h); + h.finish() + }; + let kind_name = if is_xor { "xor" } else { "range" }; + let name = format!("apc_bus_bitwise_{}_codegen_{:016x}", kind_name, source_hash); + let idx_offset_str = if is_xor { " + num_rows" } else { "" }; + + use std::fmt::Write as _; + writeln!( + s, + r#"extern "C" __global__ void {name}( + const unsigned int* __restrict__ d_output, + int N, + unsigned long long H, + unsigned int* __restrict__ d_bitwise_hist) +{{ + const int warp = (threadIdx.x >> 5); + const int lane = (threadIdx.x & 31); + const int wpb = (blockDim.x >> 5); + constexpr int N_OPS = {n_ops}; + const unsigned int num_rows = 65536u; + const unsigned int total_bins = 2u * num_rows; + lookup::Histogram hist(d_bitwise_hist, total_bins); + + for (int i = blockIdx.x * wpb + warp; i < N_OPS; i += gridDim.x * wpb) {{ + switch (i) {{"#, + ) + .unwrap(); + + let mut idx = 0usize; + for op in affine_ops { + writeln!(s, " case {}: {{", idx).unwrap(); + s.push_str(" for (int r = lane; r < N; r += 32) {\n"); + if op.guard_col != NO_GUARD { + writeln!( + s, + " unsigned int g = monty_reduce(d_output[{}ull * H + (unsigned long long)r]); + if (g == 0u) continue;", + op.guard_col + ) + .unwrap(); + } + emit_affine_inline(&mut s, &op.x, " ", "x"); + emit_affine_inline(&mut s, &op.y, " ", "y"); + writeln!( + s, + " unsigned int idx = ((x << 8) | (y & 0xFFu)){idx_offset_str};" + ) + .unwrap(); + if op.mult_const == 1 { + s.push_str(" hist.add_count(idx);\n"); + } else { + writeln!( + s, + " hist.add_count_n(idx, {}u);", + op.mult_const + ) + .unwrap(); + } + s.push_str(" }\n break;\n }\n"); + idx += 1; + } + for (_orig_i, op) in bilinear_ops { + writeln!(s, " case {}: {{ /* bilinear */", idx).unwrap(); + s.push_str(" for (int r = lane; r < N; r += 32) {\n"); + if op.guard_col != NO_GUARD { + writeln!( + s, + " unsigned int g = monty_reduce(d_output[{}ull * H + (unsigned long long)r]); + if (g == 0u) continue;", + op.guard_col + ) + .unwrap(); + } + emit_bilinear_inline(&mut s, &op.x, " ", "x"); + emit_bilinear_inline(&mut s, &op.y, " ", "y"); + writeln!( + s, + " unsigned int idx = ((x << 8) | (y & 0xFFu)){idx_offset_str};" + ) + .unwrap(); + if op.mult_const == 1 { + s.push_str(" hist.add_count(idx);\n"); + } else { + writeln!( + s, + " hist.add_count_n(idx, {}u);", + op.mult_const + ) + .unwrap(); + } + s.push_str(" }\n break;\n }\n"); + idx += 1; + } + s.push_str(" default: break;\n }\n }\n}\n"); + + Some(EmittedKernel { + source: s, + name, + source_hash, + }) +} + +/// Emit straight-line C++ that computes the affine arg into `out_var` +/// (canonical u32). Constants are baked as immediates; per-term work is +/// fully unrolled at codegen time so ptxas can FMA-fuse and constant-fold. +fn emit_affine_inline(s: &mut String, a: &AffineArg, indent: &str, out_var: &str) { + use std::fmt::Write as _; + writeln!( + s, + "{indent}unsigned int {out_var}_monty = {coef_const}u; /* monty(coef_const) */", + coef_const = a.coef_const_monty + ) + .unwrap(); + for t in 0..a.n_terms as usize { + writeln!( + s, + "{indent}{out_var}_monty = add_monty({out_var}_monty, mul_monty({coef}u, d_output[{col}ull * H + (unsigned long long)r]));", + coef = a.coefs_monty[t], + col = a.cols[t] + ) + .unwrap(); + } + writeln!(s, "{indent}unsigned int {out_var} = monty_reduce({out_var}_monty);").unwrap(); +} + +/// Emit straight-line C++ that computes a bilinear (degree-≤2) arg into +/// `out_var` (canonical u32). Constant + linear terms emitted same as +/// affine; bilinear terms add `mul_monty(c, mul_monty(cell_a, cell_b))` +/// per cross-term. Cell loads are deduped within this single emission +/// so a column referenced in both linear and bilinear positions only +/// emits one LDG. +fn emit_bilinear_inline( + s: &mut String, + m: &BilinearMonomials, + indent: &str, + out_var: &str, +) { + use std::fmt::Write as _; + // Collect all distinct columns referenced (linear + bilinear) and emit + // one LDG per column up front. ptxas can hoist these out of the row + // loop where applicable. + let mut cols: Vec = Vec::new(); + let mut add_col = |c: u32, cols: &mut Vec| { + if !cols.contains(&c) { + cols.push(c); + } + }; + for (c, _) in &m.linear { + add_col(*c, &mut cols); + } + for (a, b, _) in &m.bilinear { + add_col(*a, &mut cols); + add_col(*b, &mut cols); + } + let cell_var = |col: u32| format!("{out_var}_c{col}"); + for col in &cols { + writeln!( + s, + "{indent}unsigned int {cell} = d_output[{col}ull * H + (unsigned long long)r];", + cell = cell_var(*col) + ) + .unwrap(); + } + // Initialize accumulator with the constant in Monty form. + writeln!( + s, + "{indent}unsigned int {out_var}_monty = {const_monty}u;", + const_monty = host_to_monty(m.constant) + ) + .unwrap(); + // Linear terms. + for (col, coef) in &m.linear { + let coef_monty = host_to_monty(*coef); + writeln!( + s, + "{indent}{out_var}_monty = add_monty({out_var}_monty, mul_monty({coef_monty}u, {cell}));", + cell = cell_var(*col) + ) + .unwrap(); + } + // Bilinear terms: cross-product mul_monty followed by coef multiply. + for (col_a, col_b, coef) in &m.bilinear { + let coef_monty = host_to_monty(*coef); + if col_a == col_b { + writeln!( + s, + "{indent}{{ unsigned int p = mul_monty({cell_a}, {cell_a}); {out_var}_monty = add_monty({out_var}_monty, mul_monty({coef_monty}u, p)); }}", + cell_a = cell_var(*col_a) + ) + .unwrap(); + } else { + writeln!( + s, + "{indent}{{ unsigned int p = mul_monty({cell_a}, {cell_b}); {out_var}_monty = add_monty({out_var}_monty, mul_monty({coef_monty}u, p)); }}", + cell_a = cell_var(*col_a), + cell_b = cell_var(*col_b) + ) + .unwrap(); + } + } + writeln!(s, "{indent}unsigned int {out_var} = monty_reduce({out_var}_monty);").unwrap(); +} + +fn hash_codegen_kernel(tag: &str, ops: &[T]) -> u64 { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut hasher = DefaultHasher::new(); + "BUS_CODEGEN".hash(&mut hasher); + EMITTER_VERSION.hash(&mut hasher); + tag.hash(&mut hasher); + ops.hash(&mut hasher); + hasher.finish() +} + +/// Host-side Montgomery encode for BabyBear: `x * R mod P`. R = 2^32, P = 0x78000001. +/// Used by tests that need to load synthetic data into a column-major Monty trace. +pub fn host_to_monty(x: u32) -> u32 { + const R2: u64 = 1_172_168_163; + + fn monty_reduce(x: u64) -> u32 { + const P: u64 = 0x7800_0001; + const M_INV: u64 = 0x8800_0001; + let t = (x.wrapping_mul(M_INV)) & 0xFFFF_FFFF; + let u = t * P; + let (x_sub_u, overflow) = x.overflowing_sub(u); + let hi = (x_sub_u >> 32) as u32; + if overflow { + hi.wrapping_add(P as u32) + } else { + hi + } + } + + monty_reduce((x as u64) * R2) +} + +/// Shared device helpers + warp-dedup histogram. The `add_count_n` variant +/// applies multiplicity in a single atomicAdd so mult > 1 doesn't change +/// the dedup-mask semantics (unlike a `for k=0..mult` loop on Volta+ where +/// per-lane mult divergence corrupts `__activemask`). +const BUS_PRELUDE: &str = r#"// Auto-generated by powdr nvrtc_bus_emit. Do not edit. +// +// Self-contained — no external includes. Reads BabyBear Monty-form trace +// cells and updates `__constant__` op-table-driven periphery histograms. + +namespace lookup { +struct Histogram { + unsigned int* global_hist; + unsigned int num_bins; + __device__ __forceinline__ Histogram(unsigned int* h, unsigned int n) + : global_hist(h), num_bins(n) {} + + /// Original add_count: increment by 1 with warp-dedup. + __device__ __forceinline__ void add_count(unsigned int idx) { + if (idx < num_bins) { + unsigned int curr_mask = __activemask(); + unsigned int same_mask = __match_any_sync(curr_mask, idx); + int leader = __ffs(same_mask) - 1; + if (((int)threadIdx.x & 31) == leader) { + atomicAdd(&global_hist[idx], (unsigned int)__popc(same_mask)); + } + } + } + + /// add_count_n: increment by `mult` with warp-dedup. Required for mult > + /// 1 because the natural `for k=0..mult: add_count(idx)` diverges on + /// per-lane mult, corrupting the `__activemask` semantics on Volta+. + /// Caller must guarantee `mult` is uniform across the warp (host-side + /// it's a `__constant__` field, so this holds). + __device__ __forceinline__ void add_count_n(unsigned int idx, unsigned int mult) { + if (idx < num_bins) { + unsigned int curr_mask = __activemask(); + unsigned int same_mask = __match_any_sync(curr_mask, idx); + int leader = __ffs(same_mask) - 1; + if (((int)threadIdx.x & 31) == leader) { + atomicAdd(&global_hist[idx], mult * (unsigned int)__popc(same_mask)); + } + } + } +}; +} // namespace lookup + +__device__ __forceinline__ unsigned int monty_reduce(unsigned long long x) { + constexpr unsigned int M = 0x88000001u; + constexpr unsigned int P = 0x78000001u; + unsigned long long t = (x * (unsigned long long)M) & 0xFFFFFFFFull; + unsigned long long u = t * (unsigned long long)P; + unsigned long long x_sub_u = x - u; + bool overflow = x < u; + unsigned int hi = (unsigned int)(x_sub_u >> 32); + return hi + (overflow ? P : 0u); +} + +// Op structs match the Rust #[repr(C)] layout in nvrtc_bus_emit.rs. +// MAX_TERMS_PER_ARG = 5 here in C++, must stay in lock-step with the Rust const. +#define MAX_TERMS_PER_ARG 5 + +struct AffineArg { + unsigned int n_terms; + unsigned int cols[MAX_TERMS_PER_ARG]; + unsigned int coefs_monty[MAX_TERMS_PER_ARG]; + unsigned int coef_const_monty; +}; +struct VarRangeOp { + unsigned int mult_const; + unsigned int guard_col; + AffineArg value; + unsigned int max_bits; +}; +struct Tuple2Op { + unsigned int mult_const; + unsigned int guard_col; + AffineArg v0; + AffineArg v1; +}; +struct BitwiseOp { + unsigned int mult_const; + unsigned int guard_col; + AffineArg x; + AffineArg y; +}; + +#define NO_GUARD 0xFFFFFFFFu + +// Field arithmetic in Montgomery form. Inputs/outputs are u32 monty values. +__device__ __forceinline__ unsigned int add_monty(unsigned int a, unsigned int b) { + constexpr unsigned int Pp = 0x78000001u; + unsigned int s = a + b; + return s >= Pp ? s - Pp : s; +} +__device__ __forceinline__ unsigned int mul_monty(unsigned int a, unsigned int b) { + return monty_reduce((unsigned long long)a * (unsigned long long)b); +} + +// Evaluate `coef_const + sum_{i> 5); + const int lane = (threadIdx.x & 31); + const int wpb = (blockDim.x >> 5); + + for (int i = blockIdx.x * wpb + warp; i < (int)n_var_ops; i += gridDim.x * wpb) { + VarRangeOp op = d_ops[i]; + unsigned int bits_one = 1u << op.max_bits; + lookup::Histogram hist(d_var_hist, var_num_bins); + for (int r = lane; r < N; r += 32) { + unsigned int m = eval_mult(op.mult_const, op.guard_col, d_output, H, r); + if (m == 0u) continue; + unsigned int v = eval_affine_arg(op.value, d_output, H, r); + unsigned int idx = bits_one + v - 1u; + hist.add_count_n(idx, m); + } + } +} +"#; + +const TUPLE2_KERNEL_SRC: &str = r#" +extern "C" __global__ void apc_bus_tuple2( + const unsigned int* __restrict__ d_output, + int N, + unsigned long long H, + unsigned int* __restrict__ d_tuple2_hist, + unsigned int tuple2_sz0, + unsigned int tuple2_sz1, + unsigned int n_tup_ops, + const Tuple2Op* __restrict__ d_ops) +{ + const int warp = (threadIdx.x >> 5); + const int lane = (threadIdx.x & 31); + const int wpb = (blockDim.x >> 5); + const unsigned int total_bins = tuple2_sz0 * tuple2_sz1; + + for (int i = blockIdx.x * wpb + warp; i < (int)n_tup_ops; i += gridDim.x * wpb) { + Tuple2Op op = d_ops[i]; + lookup::Histogram hist(d_tuple2_hist, total_bins); + for (int r = lane; r < N; r += 32) { + unsigned int m = eval_mult(op.mult_const, op.guard_col, d_output, H, r); + if (m == 0u) continue; + unsigned int v0 = eval_affine_arg(op.v0, d_output, H, r); + unsigned int v1 = eval_affine_arg(op.v1, d_output, H, r); + unsigned int idx = v0 * tuple2_sz1 + v1; + hist.add_count_n(idx, m); + } + } +} +"#; + +/// Bitwise range half: idx in [0, num_rows). Shares structure with +/// var_range/tuple but with `BITWISE_NUM_BITS = 8` baked in (matches openvm). +const BITWISE_RANGE_KERNEL_SRC: &str = r#" +extern "C" __global__ void apc_bus_bitwise_range( + const unsigned int* __restrict__ d_output, + int N, + unsigned long long H, + unsigned int* __restrict__ d_bitwise_hist, + unsigned int /* unused_extra0 — kept for v2 launcher uniformity */, + unsigned int n_bit_range_ops, + const BitwiseOp* __restrict__ d_ops) +{ + const int warp = (threadIdx.x >> 5); + const int lane = (threadIdx.x & 31); + const int wpb = (blockDim.x >> 5); + const unsigned int num_rows = 65536u; /* 2^16 for BITWISE_NUM_BITS=8 */ + const unsigned int total_bins = 2u * num_rows; + + for (int i = blockIdx.x * wpb + warp; i < (int)n_bit_range_ops; i += gridDim.x * wpb) { + BitwiseOp op = d_ops[i]; + lookup::Histogram hist(d_bitwise_hist, total_bins); + for (int r = lane; r < N; r += 32) { + unsigned int m = eval_mult(op.mult_const, op.guard_col, d_output, H, r); + if (m == 0u) continue; + unsigned int x = eval_affine_arg(op.x, d_output, H, r); + unsigned int y = eval_affine_arg(op.y, d_output, H, r); + unsigned int idx = (x << 8) | (y & 0xFFu); + hist.add_count_n(idx, m); + } + } +} +"#; + +const BITWISE_XOR_KERNEL_SRC: &str = r#" +extern "C" __global__ void apc_bus_bitwise_xor( + const unsigned int* __restrict__ d_output, + int N, + unsigned long long H, + unsigned int* __restrict__ d_bitwise_hist, + unsigned int /* unused_extra0 */, + unsigned int n_bit_xor_ops, + const BitwiseOp* __restrict__ d_ops) +{ + const int warp = (threadIdx.x >> 5); + const int lane = (threadIdx.x & 31); + const int wpb = (blockDim.x >> 5); + const unsigned int num_rows = 65536u; + const unsigned int total_bins = 2u * num_rows; + + for (int i = blockIdx.x * wpb + warp; i < (int)n_bit_xor_ops; i += gridDim.x * wpb) { + BitwiseOp op = d_ops[i]; + lookup::Histogram hist(d_bitwise_hist, total_bins); + for (int r = lane; r < N; r += 32) { + unsigned int m = eval_mult(op.mult_const, op.guard_col, d_output, H, r); + if (m == 0u) continue; + unsigned int x = eval_affine_arg(op.x, d_output, H, r); + unsigned int y = eval_affine_arg(op.y, d_output, H, r); + unsigned int idx = ((x << 8) | (y & 0xFFu)) + num_rows; + hist.add_count_n(idx, m); + } + } +} +"#; + +#[cfg(test)] +mod tests { + use super::*; + use openvm_stark_backend::p3_field::PrimeCharacteristicRing; + use powdr_autoprecompiles::expression::AlgebraicReference; + + fn ref_expr(id: u64, name: &str) -> AlgebraicExpression { + AlgebraicExpression::Reference(AlgebraicReference { + id, + name: std::sync::Arc::new(name.to_string()), + }) + } + + fn num_expr(c: u32) -> AlgebraicExpression { + AlgebraicExpression::Number(BabyBear::from_u32(c)) + } + + #[test] + fn host_to_monty_known_values() { + assert_eq!(host_to_monty(0), 0); + let r_mod_p = ((1u64 << 32) % 0x7800_0001) as u32; + assert_eq!(host_to_monty(1), r_mod_p); + } + + #[test] + fn partition_simple_var_range() { + let mut id_map = BTreeMap::new(); + id_map.insert(100, 0); + let input = BusEmitterInput { + apc_height: 8, + apc_poly_id_to_index: id_map, + interactions: vec![BusInteractionDesc { + kind: BusKind::VarRange, + mult: num_expr(1), + args: vec![ref_expr(100, "v"), num_expr(4)], + }], + }; + let p = partition_apc_bus(&input).unwrap(); + assert_eq!(p.var_ops.len(), 1); + assert_eq!(p.var_ops[0].mult_const, 1); + assert_eq!(p.var_ops[0].guard_col, NO_GUARD); + assert_eq!(p.var_ops[0].value.n_terms, 1); + assert_eq!(p.var_ops[0].value.cols[0], 0); + assert_eq!(p.var_ops[0].max_bits, 4); + } + + /// Multi-term affine: `15360*col_a + 15360*col_b - 15360*col_c + 15360` + /// (the dominant keccak unhandled shape). + #[test] + fn partition_var_range_multi_term_affine() { + let mut id_map = BTreeMap::new(); + id_map.insert(100, 5); + id_map.insert(101, 6); + id_map.insert(102, 7); + id_map.insert(200, 99); // is_valid + let m = |a, b, op| { + AlgebraicExpression::BinaryOperation(powdr_expression::AlgebraicBinaryOperation { + left: Box::new(a), + op, + right: Box::new(b), + }) + }; + let mul = |a, b| m(a, b, powdr_expression::AlgebraicBinaryOperator::Mul); + let add = |a, b| m(a, b, powdr_expression::AlgebraicBinaryOperator::Add); + let sub = |a, b| m(a, b, powdr_expression::AlgebraicBinaryOperator::Sub); + // value = 15360*col_a + 15360*col_b + 15360 - 15360*col_c + let value = sub( + add( + add( + mul(num_expr(15360), ref_expr(100, "a")), + mul(num_expr(15360), ref_expr(101, "b")), + ), + num_expr(15360), + ), + mul(num_expr(15360), ref_expr(102, "c")), + ); + let input = BusEmitterInput { + apc_height: 8, + apc_poly_id_to_index: id_map, + interactions: vec![BusInteractionDesc { + kind: BusKind::VarRange, + mult: ref_expr(200, "is_valid"), + args: vec![value, num_expr(12)], + }], + }; + let p = partition_apc_bus(&input).unwrap(); + assert!(p.unhandled.is_empty(), "should be fully handled"); + assert_eq!(p.var_ops.len(), 1); + let op = &p.var_ops[0]; + assert_eq!(op.mult_const, 1); + assert_eq!(op.guard_col, 99); + assert_eq!(op.value.n_terms, 3); // 3 distinct columns + assert_eq!(op.value.coef_const_monty, host_to_monty(15360)); + // Verify the columns and coefs (sorted by appearance) + let cs: Vec = (0..3).map(|i| op.value.cols[i as usize]).collect(); + assert!(cs.contains(&5)); + assert!(cs.contains(&6)); + assert!(cs.contains(&7)); + } + + #[test] + fn partition_var_range_with_guard_column() { + let mut id_map = BTreeMap::new(); + id_map.insert(100, 0); // value + id_map.insert(200, 7); // is_valid + let input = BusEmitterInput { + apc_height: 8, + apc_poly_id_to_index: id_map, + interactions: vec![BusInteractionDesc { + kind: BusKind::VarRange, + // mult = is_valid * 1 → decode_mult collapses to is_valid (Reference) → guard_col = 7 + mult: AlgebraicExpression::BinaryOperation( + powdr_expression::AlgebraicBinaryOperation { + left: Box::new(ref_expr(200, "is_valid")), + op: powdr_expression::AlgebraicBinaryOperator::Mul, + right: Box::new(num_expr(1)), + }, + ), + args: vec![ref_expr(100, "v"), num_expr(4)], + }], + }; + let p = partition_apc_bus(&input).unwrap(); + assert_eq!(p.var_ops.len(), 1); + assert_eq!(p.var_ops[0].mult_const, 1); + assert_eq!(p.var_ops[0].guard_col, 7); + assert_eq!(p.var_ops[0].value.n_terms, 1); + assert_eq!(p.var_ops[0].value.cols[0], 0); + } + + #[test] + fn partition_collects_non_simple_as_unhandled() { + let mut id_map = BTreeMap::new(); + id_map.insert(100, 0); + let input = BusEmitterInput { + apc_height: 8, + apc_poly_id_to_index: id_map, + interactions: vec![BusInteractionDesc { + kind: BusKind::VarRange, + // mult is an Add — neither Number, Reference, nor const*ref → + // lands in `unhandled` for bytecode-VM fallback. + mult: AlgebraicExpression::BinaryOperation( + powdr_expression::AlgebraicBinaryOperation { + left: Box::new(ref_expr(100, "v")), + op: powdr_expression::AlgebraicBinaryOperator::Add, + right: Box::new(num_expr(1)), + }, + ), + args: vec![ref_expr(100, "v"), num_expr(4)], + }], + }; + let p = partition_apc_bus(&input).unwrap(); + assert!(p.var_ops.is_empty()); + assert_eq!(p.unhandled, vec![0]); + } + + #[test] + fn partition_bitwise_splits_range_xor() { + let mut id_map = BTreeMap::new(); + id_map.insert(100, 0); + id_map.insert(101, 1); + let bw = |sel: u32| BusInteractionDesc { + kind: BusKind::Bitwise, + mult: num_expr(1), + args: vec![ + ref_expr(100, "x"), + ref_expr(101, "y"), + num_expr(0), + num_expr(sel), + ], + }; + let input = BusEmitterInput { + apc_height: 8, + apc_poly_id_to_index: id_map, + interactions: vec![bw(0), bw(0), bw(1)], + }; + let p = partition_apc_bus(&input).unwrap(); + assert_eq!(p.bitwise_range_ops.len(), 2); + assert_eq!(p.bitwise_xor_ops.len(), 1); + } + + #[test] + fn kernel_sources_compile_via_nvrtc() { + use crate::powdr_extension::trace_generator::cuda::nvrtc_cache::NvrtcKernelCache; + + let cache = NvrtcKernelCache::default(); + for (sym_name, kernel) in kernel_sources() { + let compiled = cache + .get_or_compile(&kernel) + .unwrap_or_else(|e| panic!("compile {} failed: {:?}", kernel.name, e)); + assert!(!compiled.function().is_null()); + assert!(!sym_name.is_empty()); // sanity + } + } +} diff --git a/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_cache.rs b/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_cache.rs new file mode 100644 index 0000000000..49e5551d83 --- /dev/null +++ b/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_cache.rs @@ -0,0 +1,525 @@ +//! Phase 1 NVRTC kernel cache: dedupes compilation by source hash and +//! holds the resulting `CUmodule` / `CUfunction` for reuse across launches. +//! +//! No disk persistence, no warm-from-disk path — everything lives in +//! process memory. Replace with a richer cache in a later phase. + +use std::collections::HashMap; +use std::ffi::{c_void, CString}; +use std::sync::{Arc, Mutex, OnceLock}; + +use rayon::prelude::*; + +use crate::cuda_abi; + +use super::nvrtc_emit::EmittedKernel; + +/// A compiled NVRTC kernel held by the cache. Owns the loaded `CUmodule` +/// and exposes the `CUfunction` to launch. +#[derive(Debug)] +pub struct CompiledKernel { + /// Opaque `CUmodule` handle. Released on drop. + module: *mut c_void, + /// Opaque `CUfunction` handle into `module`. + function: *mut c_void, + /// Symbol name we looked up. + pub name: String, +} + +// CUDA driver handles are thread-safe to share once obtained. +unsafe impl Send for CompiledKernel {} +unsafe impl Sync for CompiledKernel {} + +impl CompiledKernel { + /// Raw `CUfunction` for use with `powdr_nvrtc_launch_jit_v1`. + pub fn function(&self) -> *mut c_void { + self.function + } + + /// Raw `CUmodule`. Used to look up `__constant__` symbols on the + /// kernel's module via `cuModuleGetGlobal`. + pub fn module(&self) -> *mut c_void { + self.module + } +} + +impl Drop for CompiledKernel { + fn drop(&mut self) { + if !self.module.is_null() { + // Best-effort unload; ignore any error during shutdown. + unsafe { + let _ = cuda_abi::powdr_nvrtc_unload_module(self.module); + } + self.module = std::ptr::null_mut(); + self.function = std::ptr::null_mut(); + } + } +} + +/// Errors from compiling and loading an NVRTC kernel. +#[derive(Debug)] +pub enum NvrtcError { + Compile { code: i32, log: Option }, + Load { code: i32 }, + GetFunction { code: i32 }, + InvalidSource, +} + +impl std::fmt::Display for NvrtcError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + NvrtcError::Compile { code, log } => { + write!(f, "NVRTC compile failed (code {})", code)?; + if let Some(l) = log { + write!(f, ": {}", l)?; + } + Ok(()) + } + NvrtcError::Load { code } => write!(f, "cuModuleLoadData failed (code {})", code), + NvrtcError::GetFunction { code } => { + write!(f, "cuModuleGetFunction failed (code {})", code) + } + NvrtcError::InvalidSource => write!(f, "invalid kernel source (interior NUL)"), + } + } +} + +impl std::error::Error for NvrtcError {} + +/// Process-wide cache of compiled NVRTC kernels. +#[derive(Default)] +pub struct NvrtcKernelCache { + by_hash: Mutex>>, +} + +impl NvrtcKernelCache { + pub fn global() -> &'static Self { + static CACHE: OnceLock = OnceLock::new(); + CACHE.get_or_init(NvrtcKernelCache::default) + } + + /// Compile `kernel` if not already cached and return a handle that can be + /// launched many times. Identical sources (matched by `source_hash`) share + /// one compiled module. + pub fn get_or_compile(&self, kernel: &EmittedKernel) -> Result, NvrtcError> { + if let Some(existing) = self.by_hash.lock().unwrap().get(&kernel.source_hash) { + return Ok(Arc::clone(existing)); + } + let compiled = compile_one(kernel)?; + let arc = Arc::new(compiled); + self.by_hash + .lock() + .unwrap() + .entry(kernel.source_hash) + .or_insert_with(|| Arc::clone(&arc)); + Ok(arc) + } + + /// Bulk parallel compile: maps `kernels[i]` → result `Arc`, + /// running cache misses concurrently across the rayon thread pool. Cache + /// hits return the existing `Arc` without recompiling. Order is preserved. + /// Duplicate hashes within `kernels` compile only once and share an `Arc`. + /// + /// Both the in-memory cache (this struct's `by_hash`) and the on-disk PTX + /// cache (in `compile_one`) are populated as a side effect, so subsequent + /// `get_or_compile` calls for the same hash return immediately. + /// + /// Useful when many APCs need first-time compilation: an APC=30 prove + /// needs ~30 NVRTC compiles cold; serially that's seconds, in parallel + /// (NVRTC + nvcc are thread-safe per-program) it's near `max(per-kernel)`. + pub fn get_or_compile_many( + &self, + kernels: &[EmittedKernel], + ) -> Vec, NvrtcError>> { + // Collect unique kernels (by source_hash) that aren't yet cached. The + // first occurrence of each hash carries the source we'll compile. + let unique_misses: Vec<&EmittedKernel> = { + let by_hash = self.by_hash.lock().unwrap(); + let mut seen: std::collections::HashSet = std::collections::HashSet::new(); + let mut keep = Vec::new(); + for k in kernels { + if by_hash.contains_key(&k.source_hash) { + continue; + } + if seen.insert(k.source_hash) { + keep.push(k); + } + } + keep + }; + + // Compile the unique misses in parallel. + let compiled: Vec<(u64, Result, NvrtcError>)> = unique_misses + .par_iter() + .map(|k| (k.source_hash, compile_one(k).map(Arc::new))) + .collect(); + + // Insert successful compiles into the cache. + { + let mut by_hash = self.by_hash.lock().unwrap(); + for (hash, result) in &compiled { + if let Ok(arc) = result { + by_hash.entry(*hash).or_insert_with(|| Arc::clone(arc)); + } + } + } + + // Map each input kernel to its result by re-reading the (now-populated) + // cache, falling back to per-hash compile errors that didn't make it + // into the cache. + let by_hash = self.by_hash.lock().unwrap(); + let mut errors: std::collections::HashMap = std::collections::HashMap::new(); + for (hash, result) in compiled { + if let Err(e) = result { + errors.insert(hash, e); + } + } + kernels + .iter() + .map(|k| { + if let Some(existing) = by_hash.get(&k.source_hash) { + Ok(Arc::clone(existing)) + } else if let Some(e) = errors.remove(&k.source_hash) { + Err(e) + } else { + // Hash missing from both cache and error map: a duplicate + // entry whose error we already returned above. Reconstruct + // an InvalidSource so callers see *some* error. + Err(NvrtcError::InvalidSource) + } + }) + .collect() + } +} + +/// Disk PTX cache directory. Honors `POWDR_NVRTC_CACHE_DIR` if set; otherwise +/// falls back to `~/.cache/powdr/nvrtc_kernels`. Returns `None` if neither +/// can be resolved (in which case the on-disk cache is silently skipped). +fn disk_cache_dir() -> Option { + if let Ok(d) = std::env::var("POWDR_NVRTC_CACHE_DIR") { + return Some(std::path::PathBuf::from(d)); + } + let home = std::env::var("HOME").ok()?; + Some(std::path::PathBuf::from(home).join(".cache/powdr/nvrtc_kernels")) +} + +fn disk_cache_path(kernel: &EmittedKernel) -> Option { + Some(disk_cache_dir()?.join(format!("{:016x}.ptx", kernel.source_hash))) +} + +fn compile_one(kernel: &EmittedKernel) -> Result { + if let Ok(dir) = std::env::var("POWDR_NVRTC_DUMP_DIR") { + let path = format!("{}/{}.cu", dir, kernel.name); + if let Err(e) = std::fs::write(&path, kernel.source.as_bytes()) { + tracing::warn!("POWDR_NVRTC_DUMP_DIR: failed to write {}: {}", path, e); + } else { + tracing::info!( + "Dumped NVRTC source ({} bytes, {} lines) to {}", + kernel.source.len(), + kernel.source.lines().count(), + path + ); + } + } + + // Disk cache fast path: load PTX bytes and skip NVRTC entirely. + if std::env::var("POWDR_NVRTC_NO_DISK_CACHE").is_err() { + if let Some(p) = disk_cache_path(kernel) { + if let Ok(ptx) = std::fs::read(&p) { + let load_start = std::time::Instant::now(); + match load_module_from_ptx(&ptx, &kernel.name) { + Ok(handle) => { + tracing::info!( + "NVRTC PTX cache hit: loaded {} from {} in {:.2}ms ({} bytes)", + kernel.name, + p.display(), + load_start.elapsed().as_secs_f64() * 1000.0, + ptx.len() + ); + return Ok(handle); + } + Err(e) => { + tracing::warn!( + "NVRTC PTX cache hit but load failed ({}); recompiling", + e + ); + // Fall through to NVRTC compile. + } + } + } + } + } + + let compile_start = std::time::Instant::now(); + let src = CString::new(kernel.source.as_str()).map_err(|_| NvrtcError::InvalidSource)?; + let src_name = CString::new("apc_kernel.cu").unwrap(); + + let mut ptx_ptr: *mut std::ffi::c_char = std::ptr::null_mut(); + let mut ptx_size: usize = 0; + let mut log: *mut std::ffi::c_char = std::ptr::null_mut(); + + let rc = unsafe { + cuda_abi::powdr_nvrtc_compile( + src.as_ptr(), + src_name.as_ptr(), + &mut ptx_ptr, + &mut ptx_size, + &mut log, + ) + }; + if rc != 0 { + let log_str = unsafe { take_c_str(log) }; + return Err(NvrtcError::Compile { code: rc, log: log_str }); + } + debug_assert!(!ptx_ptr.is_null() && ptx_size > 0, "NVRTC returned empty PTX"); + if !log.is_null() { + unsafe { cuda_abi::powdr_nvrtc_free(log) }; + } + + // Copy PTX bytes out of the NVRTC-allocated buffer so we can both load + // the module and persist to disk after the NVRTC buffer is freed. + let ptx_bytes: Vec = unsafe { + std::slice::from_raw_parts(ptx_ptr as *const u8, ptx_size).to_vec() + }; + unsafe { cuda_abi::powdr_nvrtc_free(ptx_ptr) }; + + let elapsed = compile_start.elapsed(); + tracing::info!( + "NVRTC compiled {} in {:.2}s (source {} bytes, ptx {} bytes)", + kernel.name, + elapsed.as_secs_f64(), + kernel.source.len(), + ptx_bytes.len() + ); + + let handle = load_module_from_ptx(&ptx_bytes, &kernel.name)?; + + // Best-effort persist to disk for future runs. + if std::env::var("POWDR_NVRTC_NO_DISK_CACHE").is_err() { + if let Some(p) = disk_cache_path(kernel) { + if let Some(parent) = p.parent() { + let _ = std::fs::create_dir_all(parent); + } + // Atomic write: write to a temp file in the same dir, then rename. + let tmp = p.with_extension("ptx.tmp"); + if std::fs::write(&tmp, &ptx_bytes) + .and_then(|_| std::fs::rename(&tmp, &p)) + .is_err() + { + let _ = std::fs::remove_file(&tmp); + tracing::debug!( + "Failed to persist PTX to {}; future runs will recompile", + p.display() + ); + } + } + } + + Ok(handle) +} + +/// Load PTX bytes into a CUDA module and look up the kernel function. +fn load_module_from_ptx(ptx: &[u8], kernel_name: &str) -> Result { + let name_c = CString::new(kernel_name).map_err(|_| NvrtcError::InvalidSource)?; + + let mut module: *mut c_void = std::ptr::null_mut(); + let load_rc = unsafe { + cuda_abi::powdr_nvrtc_load_module( + ptx.as_ptr() as *const c_void, + ptx.len(), + &mut module, + ) + }; + if load_rc != 0 { + return Err(NvrtcError::Load { code: load_rc }); + } + + let mut function: *mut c_void = std::ptr::null_mut(); + let fn_rc = unsafe { + cuda_abi::powdr_nvrtc_get_function(module, name_c.as_ptr(), &mut function) + }; + if fn_rc != 0 { + unsafe { cuda_abi::powdr_nvrtc_unload_module(module) }; + return Err(NvrtcError::GetFunction { code: fn_rc }); + } + + Ok(CompiledKernel { + module, + function, + name: kernel_name.to_string(), + }) +} + +/// Take ownership of a `*mut c_char` produced by `powdr_nvrtc_compile`, +/// returning a Rust string and freeing the underlying buffer. +unsafe fn take_c_str(p: *mut std::ffi::c_char) -> Option { + if p.is_null() { + return None; + } + let s = std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned(); + cuda_abi::powdr_nvrtc_free(p); + Some(s) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::powdr_extension::trace_generator::cuda::nvrtc_emit::{ + emit_jit_kernel_source, EmitterColumn, EmitterInput, EmitterInstruction, + }; + + fn sample() -> EmittedKernel { + emit_jit_kernel_source(&EmitterInput { + instructions: vec![EmitterInstruction { + arena_offset: 0, + record_stride: 64, + record_offset: 0, + columns: vec![EmitterColumn::DirectU32 { apc_col: 0, off: 0 }], + }], + }) + } + + #[test] + fn cache_compiles_and_dedupes() { + let cache = NvrtcKernelCache::default(); + let k1 = sample(); + let k2 = sample(); // identical source + let a = cache.get_or_compile(&k1).expect("first compile"); + let b = cache.get_or_compile(&k2).expect("second compile"); + // Same Arc => same compiled module. + assert!(Arc::ptr_eq(&a, &b), "expected dedupe by source hash"); + assert!(!a.function().is_null()); + } + + /// `get_or_compile_many` parallel-compiles distinct sources and dedupes + /// duplicates, populating the cache so a subsequent `get_or_compile` + /// returns the same `Arc` without recompiling. + #[test] + fn cache_parallel_compile_dedupes_and_populates() { + // Build N distinct sources (different stride values produce different + // hashes). Repeat one of them to test in-batch dedupe. + let mut kernels: Vec = (0..4) + .map(|i| { + emit_jit_kernel_source(&EmitterInput { + instructions: vec![EmitterInstruction { + arena_offset: 0, + record_stride: 64 + i as u32, + record_offset: 0, + columns: vec![EmitterColumn::DirectU32 { apc_col: 0, off: 0 }], + }], + }) + }) + .collect(); + // Duplicate kernels[1] so the batch contains two entries with the + // same source hash. + kernels.push(kernels[1].clone()); + + let cache = NvrtcKernelCache::default(); + let results = cache.get_or_compile_many(&kernels); + assert_eq!(results.len(), kernels.len()); + let arcs: Vec> = results + .into_iter() + .map(|r| r.expect("each compile must succeed")) + .collect(); + + // Duplicate kernel must dedupe to the same Arc as the original. + assert!( + Arc::ptr_eq(&arcs[1], &arcs[4]), + "expected duplicate kernel to dedupe within batch" + ); + + // Cache must be populated: get_or_compile returns the SAME Arc. + for (k, expected) in kernels.iter().zip(arcs.iter()) { + let got = cache.get_or_compile(k).expect("post-batch get"); + assert!(Arc::ptr_eq(&got, expected), "expected cache hit"); + } + } + + /// End-to-end Phase 1 vertical slice: build a synthetic arena with + /// known values, emit a DirectU32-only kernel, compile via NVRTC, launch + /// it against a fresh trace buffer, copy back, and assert each cell + /// matches the host-side Montgomery encoding of the synthetic value. + #[test] + fn launch_direct_u32_writes_montgomery_words() { + use openvm_cuda_common::{ + copy::{MemCopyD2H, MemCopyH2D}, + d_buffer::DeviceBuffer, + }; + + use crate::powdr_extension::trace_generator::cuda::nvrtc_emit::host_to_monty; + + // 16 valid rows, width = 3 cols. Trace is column-major. + const N: usize = 16; + const H: usize = 16; // already a power of two + const WIDTH: usize = 3; + const RECORD_STRIDE: usize = 16; // bytes per APC call + + // Build synthetic arena: row r has u32(r+1) at offset 0, u32(0xCAFE0000+r) + // at offset 4, and u32((P-1) - r) at offset 8 (last is canonical-form + // BabyBear which is what records carry). + let p: u32 = 0x7800_0001; + let mut arena_bytes = vec![0u8; N * RECORD_STRIDE]; + for r in 0..N { + let row = &mut arena_bytes[r * RECORD_STRIDE..(r + 1) * RECORD_STRIDE]; + row[0..4].copy_from_slice(&((r as u32) + 1).to_le_bytes()); + row[4..8].copy_from_slice(&(0xCAFE_0000u32 + r as u32).to_le_bytes()); + row[8..12].copy_from_slice(&((p - 1) - r as u32).to_le_bytes()); + // 12..16 left as zero + } + let d_arena = arena_bytes.to_device().expect("arena H2D"); + + // Allocate column-major trace buffer of u32 Montgomery words. Pre-zero + // so non-written cells stay 0 (Montgomery(0) = 0). + let d_trace: DeviceBuffer = DeviceBuffer::with_capacity(H * WIDTH); + d_trace.fill_zero().expect("zero trace"); + + // Emit DirectU32 kernel: col 0 <- offset 0, col 1 <- offset 4, + // col 2 <- offset 8. + let kernel = crate::powdr_extension::trace_generator::cuda::nvrtc_emit::emit_jit_kernel_source( + &EmitterInput { + instructions: vec![EmitterInstruction { + arena_offset: 0, + record_stride: RECORD_STRIDE as u32, + record_offset: 0, + columns: vec![ + EmitterColumn::DirectU32 { apc_col: 0, off: 0 }, + EmitterColumn::DirectU32 { apc_col: 1, off: 4 }, + EmitterColumn::DirectU32 { apc_col: 2, off: 8 }, + ], + }], + }, + ); + + let cache = NvrtcKernelCache::default(); + let compiled = cache.get_or_compile(&kernel).expect("compile kernel"); + + // Launch: 1 block of 256 threads is plenty for 16 rows. + let rc = unsafe { + crate::cuda_abi::powdr_nvrtc_launch_jit_v1( + compiled.function(), + d_trace.as_mut_ptr(), + H, + N as i32, + d_arena.as_ptr(), + /* range_max_bits */ 0, + /* grid_x */ 1, + /* block_x */ 256, + ) + }; + assert_eq!(rc, 0, "launch failed: {}", rc); + + let host: Vec = d_trace.to_host().expect("trace D2H"); + assert_eq!(host.len(), H * WIDTH); + + // Verify column-major layout: cell (col, row) lives at host[col*H + row]. + for r in 0..N { + let v0 = (r as u32) + 1; + let v1 = 0xCAFE_0000u32 + r as u32; + let v2 = (p - 1) - r as u32; + assert_eq!(host[0 * H + r], host_to_monty(v0), "col0 row{}", r); + assert_eq!(host[1 * H + r], host_to_monty(v1), "col1 row{}", r); + assert_eq!(host[2 * H + r], host_to_monty(v2), "col2 row{}", r); + } + // Padding rows (none here since N==H) and unused col bytes stay zero. + } +} diff --git a/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_emit.rs b/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_emit.rs new file mode 100644 index 0000000000..dd9e246b66 --- /dev/null +++ b/openvm/src/powdr_extension/trace_generator/cuda/nvrtc_emit.rs @@ -0,0 +1,1085 @@ +//! NVRTC emitter: turns surviving-column descriptors into CUDA source for a +//! per-APC trace-gen kernel. +//! +//! The host pre-folds opcodes (each APC instruction has a fixed opcode), so +//! the emitted code never dispatches on `record[opcode_byte_offset]` for +//! ALU/Shift/LoadStore — instead it emits the specific arm directly. +//! +//! Output kernel signature: +//! +//! ```text +//! extern "C" __global__ void ( +//! uint32_t* d_output, // column-major, Montgomery words +//! unsigned long long H, // trace height (power of 2) +//! int N, // number of valid rows +//! const uint8_t* d_arena, // concatenated arena bytes +//! uint32_t range_max_bits); +//! ``` + +use std::collections::hash_map::DefaultHasher; +use std::fmt::Write as _; +use std::hash::{Hash, Hasher}; + +/// Bumped whenever the emitter's output format changes — forces cache +/// invalidation across emitter revisions. +const EMITTER_VERSION: u32 = 3; + +/// One column to materialize from this instruction's record bytes. +/// +/// Each variant maps to a closed expression over `rec + offset` byte slots, +/// with all opcodes / shift directions / branch directions already folded by +/// the host (so the emitted code does not switch on opcode bytes at runtime). +#[derive(Clone, Debug, Hash)] +pub enum EmitterColumn { + /// Direct `*(uint32_t*)(rec + off)`. + DirectU32 { apc_col: u16, off: u16 }, + /// Direct `(uint32_t)rec[off]`. + DirectU8 { apc_col: u16, off: u16 }, + /// Direct `*(uint16_t*)(rec + off)`. + DirectU16 { apc_col: u16, off: u16 }, + /// Compile-time constant value (already canonical, < P). + Constant { apc_col: u16, value: u32 }, + /// `1` if the host-known instruction opcode equals `expected_opcode`, + /// else `0`. The host has already resolved the comparison so this is + /// emitted as a literal `0` or `1`. + BoolFromOpcodeFolded { apc_col: u16, value: u8 }, + /// Conditional wrapper: if `record[cond_off] != 0`, materialize `inner`, + /// else write `0`. + Conditional { + cond_off: u16, + inner: Box, + }, + /// One limb of `decompose((curr_ts + delta) - prev_ts - 1)` against + /// `range_max_bits` (kernel arg). + TimestampDecomp { + apc_col: u16, + curr_off: u16, + prev_off: u16, + delta: u16, + limb_index: u16, + }, + /// One u16 limb of `(val + sign_extend(imm, sign))`. + PointerLimb { + apc_col: u16, + val_off: u16, + imm_off: u16, + imm_sign_off: u16, + limb_index: u16, + }, + /// One byte of `run_alu(opcode, b, c)`. + AluResult { + apc_col: u16, + opcode_off: u16, + b_off: u16, + c_off: u16, + limb_index: u16, + }, + /// One byte of `run_shift(opcode, b, c)`. + ShiftResult { + apc_col: u16, + opcode_off: u16, + b_off: u16, + c_off: u16, + limb_index: u16, + }, + /// `is_sll ? (1 << bit_shift) : 0`. `is_sll` from opcode==0; bit shift from `c[0]`. + ShiftBitMulLeft { apc_col: u16, opcode_off: u16, c_off: u16 }, + /// `is_sll ? 0 : (1 << bit_shift)`. + ShiftBitMulRight { apc_col: u16, opcode_off: u16, c_off: u16 }, + /// `is_sra ? (b[3] >> 7) : 0`. `is_sra` from opcode==2. + ShiftBSign { apc_col: u16, opcode_off: u16, b_off: u16 }, + /// `bit_shift == marker_index`. + ShiftBitMarker { apc_col: u16, c_off: u16, marker_index: u16 }, + /// `limb_shift == marker_index`. + ShiftLimbMarker { apc_col: u16, c_off: u16, marker_index: u16 }, + /// Carry between limbs for SLL/SRL. See run_shift in apc_jit_tracegen.cu. + ShiftBitCarry { + apc_col: u16, + opcode_off: u16, + b_off: u16, + c_off: u16, + limb_index: u16, + }, + /// BranchEqual: cmp_result. `1` if (BEQ and a==b) or (BNE and a!=b), else 0. + BranchEqualCmpResult { + apc_col: u16, + a_off: u16, + b_off: u16, + opcode_off: u16, + }, + /// BranchEqual: diff_inv_marker[marker_index]. Field inverse of the + /// first differing limb when marker matches, else 0. + BranchEqualDiffInvMarker { + apc_col: u16, + a_off: u16, + b_off: u16, + marker_index: u16, + }, + /// LoadStore: rd_rs2_ptr (0 if 0xFFFFFFFF, else value). + LoadStoreRdRs2Ptr { apc_col: u16, off: u16 }, + /// LoadStore: needs_write (rd_rs2_ptr != 0xFFFFFFFF). + LoadStoreNeedsWrite { apc_col: u16, off: u16 }, + /// LoadStore: write_aux.prev_timestamp (gated on needs_write). + LoadStoreWriteAuxPrevTs { + apc_col: u16, + write_prev_ts_off: u16, + rd_rs2_ptr_off: u16, + }, + /// LoadStore: write_aux timestamp decomposition (gated on needs_write). + LoadStoreWriteAuxDecomp { + apc_col: u16, + from_ts_off: u16, + write_prev_ts_off: u16, + rd_rs2_ptr_off: u16, + limb_index: u16, + }, + /// LoadStore: is_load = (opcode <= 2). + LoadStoreIsLoad { apc_col: u16, opcode_off: u16 }, + /// LoadStore: flags[flag_index] from opcode + shift. + LoadStoreFlag { + apc_col: u16, + opcode_off: u16, + shift_off: u16, + flag_index: u16, + }, + /// LoadStore: write_data byte for one limb, op-dependent. + LoadStoreWriteData { + apc_col: u16, + opcode_off: u16, + shift_off: u16, + read_data_off: u16, + prev_data_off: u16, + limb_index: u16, + }, +} + +impl EmitterColumn { + /// APC column slot this emitter writes to. For `Conditional`, it's + /// the inner's apc_col. + fn apc_col(&self) -> u16 { + match self { + EmitterColumn::DirectU32 { apc_col, .. } + | EmitterColumn::DirectU8 { apc_col, .. } + | EmitterColumn::DirectU16 { apc_col, .. } + | EmitterColumn::Constant { apc_col, .. } + | EmitterColumn::BoolFromOpcodeFolded { apc_col, .. } + | EmitterColumn::TimestampDecomp { apc_col, .. } + | EmitterColumn::PointerLimb { apc_col, .. } + | EmitterColumn::AluResult { apc_col, .. } + | EmitterColumn::ShiftResult { apc_col, .. } + | EmitterColumn::ShiftBitMulLeft { apc_col, .. } + | EmitterColumn::ShiftBitMulRight { apc_col, .. } + | EmitterColumn::ShiftBSign { apc_col, .. } + | EmitterColumn::ShiftBitMarker { apc_col, .. } + | EmitterColumn::ShiftLimbMarker { apc_col, .. } + | EmitterColumn::ShiftBitCarry { apc_col, .. } + | EmitterColumn::BranchEqualCmpResult { apc_col, .. } + | EmitterColumn::BranchEqualDiffInvMarker { apc_col, .. } + | EmitterColumn::LoadStoreRdRs2Ptr { apc_col, .. } + | EmitterColumn::LoadStoreNeedsWrite { apc_col, .. } + | EmitterColumn::LoadStoreWriteAuxPrevTs { apc_col, .. } + | EmitterColumn::LoadStoreWriteAuxDecomp { apc_col, .. } + | EmitterColumn::LoadStoreIsLoad { apc_col, .. } + | EmitterColumn::LoadStoreFlag { apc_col, .. } + | EmitterColumn::LoadStoreWriteData { apc_col, .. } => *apc_col, + EmitterColumn::Conditional { inner, .. } => inner.apc_col(), + } + } +} + +/// One instruction within an APC. +#[derive(Clone, Debug, Hash)] +pub struct EmitterInstruction { + pub arena_offset: u32, + pub record_stride: u32, + pub record_offset: u32, + pub columns: Vec, +} + +/// Emitter input: per-APC-chip data the codegen needs. +#[derive(Clone, Debug, Hash)] +pub struct EmitterInput { + pub instructions: Vec, +} + +#[derive(Clone, Debug)] +pub struct EmittedKernel { + pub source: String, + pub name: String, + pub source_hash: u64, +} + +/// Emit a CUDA source string for the given `EmitterInput`. The kernel name +/// is `apc_kernel_` where `` is a hex rendering of the +/// input's structural hash. That same hash is the cache key. +pub fn emit_jit_kernel_source(input: &EmitterInput) -> EmittedKernel { + let mut hasher = DefaultHasher::new(); + EMITTER_VERSION.hash(&mut hasher); + input.hash(&mut hasher); + let source_hash = hasher.finish(); + let name = format!("apc_kernel_{:016x}", source_hash); + + let mut s = String::new(); + s.push_str(PRELUDE); + + writeln!(s, "extern \"C\" __global__ void {}(", name).unwrap(); + s.push_str(" unsigned int* __restrict__ d_output,\n"); + s.push_str(" unsigned long long H,\n"); + s.push_str(" int N,\n"); + s.push_str(" const unsigned char* __restrict__ d_arena,\n"); + s.push_str(" unsigned int range_max_bits) {\n"); + s.push_str( + " const unsigned long long total = (unsigned long long)gridDim.x * blockDim.x;\n", + ); + s.push_str( + " const unsigned long long tid = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;\n", + ); + s.push_str(" for (unsigned long long r = tid; r < (unsigned long long)N; r += total) {\n"); + + for (i, instr) in input.instructions.iter().enumerate() { + writeln!(s, " // === Instruction {} ===", i).unwrap(); + writeln!( + s, + " const unsigned char* rec{} = d_arena + {}u + {}u + r * {}u;", + i, instr.arena_offset, instr.record_offset, instr.record_stride + ) + .unwrap(); + emit_instruction_columns(&mut s, i, &instr.columns); + } + + s.push_str(" }\n"); + s.push_str("}\n"); + + EmittedKernel { + source: s, + name, + source_hash, + } +} + +/// Emit one instruction's columns with multi-limb fusion. Columns that share +/// a record signature (e.g. four `AluResult` columns differing only in +/// `limb_index`) collapse into a single `powdr_run_alu` call followed by one +/// store per limb apc_col. Non-fusable columns emit individually. +/// +/// Fusion is conservative: bare (non-Conditional) `AluResult`, `ShiftResult`, +/// and `TimestampDecomp` columns participate. Conditional-wrapped variants +/// pass through unfused because the conditional guard differs. +fn emit_instruction_columns(s: &mut String, i: usize, cols: &[EmitterColumn]) { + use std::collections::BTreeMap; + + // Fusion buckets keyed by the shared-signature tuple. Vec preserves first + // appearance order so emit order remains deterministic regardless of + // BTreeMap iteration. + let mut alu_groups: BTreeMap<(u16, u16, u16), Vec<(u16, u16)>> = BTreeMap::new(); + let mut shift_groups: BTreeMap<(u16, u16, u16), Vec<(u16, u16)>> = BTreeMap::new(); + let mut ts_groups: BTreeMap<(u16, u16, u16), Vec<(u16, u16)>> = BTreeMap::new(); + // Non-fusable columns retain original order. + let mut others: Vec<&EmitterColumn> = Vec::new(); + + for col in cols { + match col { + EmitterColumn::AluResult { apc_col, opcode_off, b_off, c_off, limb_index } => { + alu_groups + .entry((*opcode_off, *b_off, *c_off)) + .or_default() + .push((*apc_col, *limb_index)); + } + EmitterColumn::ShiftResult { apc_col, opcode_off, b_off, c_off, limb_index } => { + shift_groups + .entry((*opcode_off, *b_off, *c_off)) + .or_default() + .push((*apc_col, *limb_index)); + } + EmitterColumn::TimestampDecomp { apc_col, curr_off, prev_off, delta, limb_index } => { + ts_groups + .entry((*curr_off, *prev_off, *delta)) + .or_default() + .push((*apc_col, *limb_index)); + } + _ => others.push(col), + } + } + + // Emit fused ALU groups. + for ((opcode_off, b_off, c_off), limbs) in &alu_groups { + if limbs.len() == 1 { + let (apc_col, limb_index) = limbs[0]; + emit_column_body( + s, + i, + &EmitterColumn::AluResult { + apc_col, + opcode_off: *opcode_off, + b_off: *b_off, + c_off: *c_off, + limb_index, + }, + 8, + ); + } else { + writeln!(s, " {{").unwrap(); + writeln!(s, " unsigned char a[4];").unwrap(); + writeln!( + s, + " powdr_run_alu(rec{}[{}u], rec{} + {}u, rec{} + {}u, a);", + i, opcode_off, i, b_off, i, c_off + ).unwrap(); + for (apc_col, limb_index) in limbs { + writeln!( + s, + " d_output[(unsigned long long){}u * H + r] = to_monty((unsigned int)a[{}u]);", + apc_col, limb_index + ).unwrap(); + } + writeln!(s, " }}").unwrap(); + } + } + + // Emit fused Shift groups. + for ((opcode_off, b_off, c_off), limbs) in &shift_groups { + if limbs.len() == 1 { + let (apc_col, limb_index) = limbs[0]; + emit_column_body( + s, + i, + &EmitterColumn::ShiftResult { + apc_col, + opcode_off: *opcode_off, + b_off: *b_off, + c_off: *c_off, + limb_index, + }, + 8, + ); + } else { + writeln!(s, " {{").unwrap(); + writeln!(s, " unsigned char a[4];").unwrap(); + writeln!( + s, + " powdr_run_shift(rec{}[{}u], rec{} + {}u, rec{} + {}u, a);", + i, opcode_off, i, b_off, i, c_off + ).unwrap(); + for (apc_col, limb_index) in limbs { + writeln!( + s, + " d_output[(unsigned long long){}u * H + r] = to_monty((unsigned int)a[{}u]);", + apc_col, limb_index + ).unwrap(); + } + writeln!(s, " }}").unwrap(); + } + } + + // Emit fused TimestampDecomp groups (share curr/prev/delta -> one diff). + for ((curr_off, prev_off, delta), limbs) in &ts_groups { + if limbs.len() == 1 { + let (apc_col, limb_index) = limbs[0]; + emit_column_body( + s, + i, + &EmitterColumn::TimestampDecomp { + apc_col, + curr_off: *curr_off, + prev_off: *prev_off, + delta: *delta, + limb_index, + }, + 8, + ); + } else { + writeln!(s, " {{").unwrap(); + writeln!(s, " unsigned int curr = *(const unsigned int*)(rec{} + {}u);", i, curr_off).unwrap(); + writeln!(s, " unsigned int prev = *(const unsigned int*)(rec{} + {}u);", i, prev_off).unwrap(); + writeln!(s, " unsigned int diff = (curr + {}u) - prev - 1u;", delta).unwrap(); + writeln!(s, " unsigned int mask = (1u << range_max_bits) - 1u;", ).unwrap(); + for (apc_col, limb_index) in limbs { + writeln!( + s, + " d_output[(unsigned long long){}u * H + r] = to_monty((diff >> (range_max_bits * {}u)) & mask);", + apc_col, limb_index + ).unwrap(); + } + writeln!(s, " }}").unwrap(); + } + } + + // Non-fusable columns last (preserves original relative order among them). + for col in others { + emit_column(s, i, col); + } +} + +/// Emit the full statements for one column. Some arms need temporaries; each +/// is wrapped in its own `{ ... }` block to keep names from colliding. The +/// output buffer is pre-zeroed by the caller, so `Conditional` can emit a +/// guarded write with no explicit else branch. +fn emit_column(s: &mut String, i: usize, col: &EmitterColumn) { + match col { + EmitterColumn::Conditional { cond_off, inner } => { + writeln!( + s, + " if (rec{}[{}u] != 0u) {{", + i, cond_off + ) + .unwrap(); + emit_column_body(s, i, inner, /* indent */ 12); + writeln!(s, " }}").unwrap(); + } + _ => emit_column_body(s, i, col, /* indent */ 8), + } +} + +fn emit_column_body(s: &mut String, i: usize, col: &EmitterColumn, indent: usize) { + let pad = " ".repeat(indent); + match col { + EmitterColumn::DirectU32 { apc_col, off } => { + writeln!( + s, + "{}d_output[(unsigned long long){}u * H + r] = to_monty(*(const unsigned int*)(rec{} + {}u));", + pad, apc_col, i, off + ) + .unwrap(); + } + EmitterColumn::DirectU8 { apc_col, off } => { + writeln!( + s, + "{}d_output[(unsigned long long){}u * H + r] = to_monty((unsigned int)rec{}[{}u]);", + pad, apc_col, i, off + ) + .unwrap(); + } + EmitterColumn::DirectU16 { apc_col, off } => { + writeln!( + s, + "{}d_output[(unsigned long long){}u * H + r] = to_monty((unsigned int)(*(const unsigned short*)(rec{} + {}u)));", + pad, apc_col, i, off + ) + .unwrap(); + } + EmitterColumn::Constant { apc_col, value } => { + writeln!( + s, + "{}d_output[(unsigned long long){}u * H + r] = 0x{:08x}u; // monty({})", + pad, apc_col, host_to_monty(*value), value + ) + .unwrap(); + } + EmitterColumn::BoolFromOpcodeFolded { apc_col, value } => { + let v = *value as u32; + writeln!( + s, + "{}d_output[(unsigned long long){}u * H + r] = 0x{:08x}u; // monty({})", + pad, apc_col, host_to_monty(v), v + ) + .unwrap(); + } + EmitterColumn::TimestampDecomp { + apc_col, + curr_off, + prev_off, + delta, + limb_index, + } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} unsigned int curr = *(const unsigned int*)(rec{} + {}u);", pad, i, curr_off).unwrap(); + writeln!(s, "{} unsigned int prev = *(const unsigned int*)(rec{} + {}u);", pad, i, prev_off).unwrap(); + writeln!(s, "{} unsigned int diff = (curr + {}u) - prev - 1u;", pad, delta).unwrap(); + writeln!(s, "{} unsigned int mask = (1u << range_max_bits) - 1u;", pad).unwrap(); + writeln!(s, "{} unsigned int v = (diff >> (range_max_bits * {}u)) & mask;", pad, limb_index).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ) + .unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::PointerLimb { + apc_col, + val_off, + imm_off, + imm_sign_off, + limb_index, + } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} unsigned int val = *(const unsigned int*)(rec{} + {}u);", pad, i, val_off).unwrap(); + writeln!(s, "{} unsigned short imm = *(const unsigned short*)(rec{} + {}u);", pad, i, imm_off).unwrap(); + writeln!(s, "{} unsigned int sign = (unsigned int)rec{}[{}u];", pad, i, imm_sign_off).unwrap(); + writeln!(s, "{} unsigned int imm_ext = (unsigned int)imm + (sign ? 0xFFFF0000u : 0u);", pad).unwrap(); + writeln!(s, "{} unsigned int ptr = val + imm_ext;", pad).unwrap(); + writeln!(s, "{} unsigned int v = (ptr >> (16u * {}u)) & 0xFFFFu;", pad, limb_index).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ) + .unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::AluResult { + apc_col, + opcode_off, + b_off, + c_off, + limb_index, + } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} unsigned char a[4];", pad).unwrap(); + writeln!( + s, + "{} powdr_run_alu(rec{}[{}u], rec{} + {}u, rec{} + {}u, a);", + pad, i, opcode_off, i, b_off, i, c_off + ).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty((unsigned int)a[{}u]);", + pad, apc_col, limb_index + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::ShiftResult { + apc_col, + opcode_off, + b_off, + c_off, + limb_index, + } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} unsigned char a[4];", pad).unwrap(); + writeln!( + s, + "{} powdr_run_shift(rec{}[{}u], rec{} + {}u, rec{} + {}u, a);", + pad, i, opcode_off, i, b_off, i, c_off + ).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty((unsigned int)a[{}u]);", + pad, apc_col, limb_index + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::ShiftBitMulLeft { apc_col, opcode_off, c_off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} bool is_sll = rec{}[{}u] == 0u;", pad, i, opcode_off).unwrap(); + writeln!(s, "{} int ls, bs; powdr_get_shift_amounts(rec{}[{}u], &ls, &bs);", pad, i, c_off).unwrap(); + writeln!(s, "{} unsigned int v = is_sll ? (1u << bs) : 0u;", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::ShiftBitMulRight { apc_col, opcode_off, c_off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} bool is_sll = rec{}[{}u] == 0u;", pad, i, opcode_off).unwrap(); + writeln!(s, "{} int ls, bs; powdr_get_shift_amounts(rec{}[{}u], &ls, &bs);", pad, i, c_off).unwrap(); + writeln!(s, "{} unsigned int v = is_sll ? 0u : (1u << bs);", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::ShiftBSign { apc_col, opcode_off, b_off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} bool is_sra = rec{}[{}u] == 2u;", pad, i, opcode_off).unwrap(); + writeln!( + s, + "{} unsigned int v = is_sra ? (unsigned int)(rec{}[{}u + 3u] >> 7) : 0u;", + pad, i, b_off + ).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::ShiftBitMarker { apc_col, c_off, marker_index } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} int ls, bs; powdr_get_shift_amounts(rec{}[{}u], &ls, &bs);", pad, i, c_off).unwrap(); + writeln!(s, "{} unsigned int v = (bs == {}) ? 1u : 0u;", pad, marker_index).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::ShiftLimbMarker { apc_col, c_off, marker_index } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} int ls, bs; powdr_get_shift_amounts(rec{}[{}u], &ls, &bs);", pad, i, c_off).unwrap(); + writeln!(s, "{} unsigned int v = (ls == {}) ? 1u : 0u;", pad, marker_index).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::ShiftBitCarry { apc_col, opcode_off, b_off, c_off, limb_index } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} bool is_sll = rec{}[{}u] == 0u;", pad, i, opcode_off).unwrap(); + writeln!(s, "{} int ls, bs; powdr_get_shift_amounts(rec{}[{}u], &ls, &bs);", pad, i, c_off).unwrap(); + writeln!(s, "{} unsigned int v = 0u;", pad).unwrap(); + writeln!(s, "{} if (bs != 0) {{", pad).unwrap(); + writeln!( + s, + "{} unsigned char bv = rec{}[{}u + {}u];", + pad, i, b_off, limb_index + ).unwrap(); + writeln!( + s, + "{} unsigned char carry = is_sll ? (unsigned char)(bv >> (8 - bs)) : (unsigned char)(bv & ((1u << bs) - 1u));", + pad + ).unwrap(); + writeln!(s, "{} v = (unsigned int)carry;", pad).unwrap(); + writeln!(s, "{} }}", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::BranchEqualCmpResult { apc_col, a_off, b_off, opcode_off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} const unsigned char* aa = rec{} + {}u;", pad, i, a_off).unwrap(); + writeln!(s, "{} const unsigned char* bb = rec{} + {}u;", pad, i, b_off).unwrap(); + writeln!(s, "{} bool is_beq = rec{}[{}u] == 0u;", pad, i, opcode_off).unwrap(); + writeln!( + s, + "{} bool eq = (aa[0]==bb[0] && aa[1]==bb[1] && aa[2]==bb[2] && aa[3]==bb[3]);", + pad + ).unwrap(); + writeln!(s, "{} unsigned int v = (is_beq ? eq : !eq) ? 1u : 0u;", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::BranchEqualDiffInvMarker { apc_col, a_off, b_off, marker_index } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} const unsigned char* aa = rec{} + {}u;", pad, i, a_off).unwrap(); + writeln!(s, "{} const unsigned char* bb = rec{} + {}u;", pad, i, b_off).unwrap(); + writeln!(s, "{} int diff_idx = 4;", pad).unwrap(); + writeln!(s, "{} for (int k = 0; k < 4; ++k) {{ if (aa[k] != bb[k]) {{ diff_idx = k; break; }} }}", pad).unwrap(); + writeln!(s, "{} if (diff_idx == 4) diff_idx = 0;", pad).unwrap(); + writeln!( + s, + "{} bool all_eq = (aa[0]==bb[0] && aa[1]==bb[1] && aa[2]==bb[2] && aa[3]==bb[3]);", + pad + ).unwrap(); + writeln!(s, "{} unsigned int monty = 0u;", pad).unwrap(); + writeln!( + s, + "{} if (diff_idx == {} && !all_eq) {{", + pad, marker_index + ).unwrap(); + // av = Monty(aa[diff_idx]); bv = Monty(bb[diff_idx]); diff_monty = av - bv (mod P). + // Then monty(inv) = Monty(canonical_inv(canonical(diff_monty))). + writeln!( + s, + "{} unsigned int av = to_monty_canonical((unsigned int)aa[diff_idx]);", + pad + ).unwrap(); + writeln!( + s, + "{} unsigned int bv = to_monty_canonical((unsigned int)bb[diff_idx]);", + pad + ).unwrap(); + writeln!(s, "{} constexpr unsigned int P = 0x78000001u;", pad).unwrap(); + writeln!( + s, + "{} unsigned int diff_monty = (av >= bv) ? (av - bv) : (P - (bv - av));", + pad + ).unwrap(); + // Convert diff_monty (Monty) → canonical via monty_reduce(diff_monty * 1). + writeln!( + s, + "{} unsigned int diff_canon = monty_reduce((unsigned long long)diff_monty);", + pad + ).unwrap(); + writeln!(s, "{} unsigned int inv_canon = powdr_canonical_inv(diff_canon);", pad).unwrap(); + writeln!(s, "{} monty = to_monty_canonical(inv_canon);", pad).unwrap(); + writeln!(s, "{} }}", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = monty;", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::LoadStoreRdRs2Ptr { apc_col, off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!( + s, + "{} unsigned int v = *(const unsigned int*)(rec{} + {}u);", + pad, i, off + ).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = (v == 0xFFFFFFFFu) ? 0u : to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::LoadStoreNeedsWrite { apc_col, off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!( + s, + "{} unsigned int v = *(const unsigned int*)(rec{} + {}u);", + pad, i, off + ).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v != 0xFFFFFFFFu ? 1u : 0u);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::LoadStoreWriteAuxPrevTs { apc_col, write_prev_ts_off, rd_rs2_ptr_off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!( + s, + "{} unsigned int rd = *(const unsigned int*)(rec{} + {}u);", + pad, i, rd_rs2_ptr_off + ).unwrap(); + writeln!(s, "{} unsigned int monty = 0u;", pad).unwrap(); + writeln!(s, "{} if (rd != 0xFFFFFFFFu) {{", pad).unwrap(); + writeln!( + s, + "{} unsigned int v = *(const unsigned int*)(rec{} + {}u);", + pad, i, write_prev_ts_off + ).unwrap(); + writeln!(s, "{} monty = to_monty(v);", pad).unwrap(); + writeln!(s, "{} }}", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = monty;", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::LoadStoreWriteAuxDecomp { + apc_col, + from_ts_off, + write_prev_ts_off, + rd_rs2_ptr_off, + limb_index, + } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!( + s, + "{} unsigned int rd = *(const unsigned int*)(rec{} + {}u);", + pad, i, rd_rs2_ptr_off + ).unwrap(); + writeln!(s, "{} unsigned int monty = 0u;", pad).unwrap(); + writeln!(s, "{} if (rd != 0xFFFFFFFFu) {{", pad).unwrap(); + writeln!( + s, + "{} unsigned int curr = *(const unsigned int*)(rec{} + {}u);", + pad, i, from_ts_off + ).unwrap(); + writeln!( + s, + "{} unsigned int prev = *(const unsigned int*)(rec{} + {}u);", + pad, i, write_prev_ts_off + ).unwrap(); + writeln!(s, "{} unsigned int diff = (curr + 2u) - prev - 1u;", pad).unwrap(); + writeln!(s, "{} unsigned int mask = (1u << range_max_bits) - 1u;", pad).unwrap(); + writeln!( + s, + "{} unsigned int v = (diff >> (range_max_bits * {}u)) & mask;", + pad, limb_index + ).unwrap(); + writeln!(s, "{} monty = to_monty(v);", pad).unwrap(); + writeln!(s, "{} }}", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = monty;", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::LoadStoreIsLoad { apc_col, opcode_off } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!( + s, + "{} unsigned int v = (rec{}[{}u] <= 2u) ? 1u : 0u;", + pad, i, opcode_off + ).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(v);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::LoadStoreFlag { apc_col, opcode_off, shift_off, flag_index } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} unsigned char op = rec{}[{}u];", pad, i, opcode_off).unwrap(); + writeln!(s, "{} unsigned char sh = rec{}[{}u];", pad, i, shift_off).unwrap(); + writeln!(s, "{} unsigned char flags[4] = {{0u, 0u, 0u, 0u}};", pad).unwrap(); + writeln!(s, "{} switch (op) {{", pad).unwrap(); + writeln!(s, "{} case 0u: flags[0]=2u; break;", pad).unwrap(); + writeln!(s, "{} case 2u: if (sh==0u) flags[1]=2u; else if (sh==2u) flags[2]=2u; break;", pad).unwrap(); + writeln!(s, "{} case 1u: switch(sh){{case 0u:flags[3]=2u;break;case 1u:flags[0]=1u;break;case 2u:flags[1]=1u;break;case 3u:flags[2]=1u;break;}} break;", pad).unwrap(); + writeln!(s, "{} case 3u: flags[3]=1u; break;", pad).unwrap(); + writeln!(s, "{} case 4u: if (sh==0u){{flags[0]=1u;flags[1]=1u;}} else if (sh==2u){{flags[0]=1u;flags[2]=1u;}} break;", pad).unwrap(); + writeln!(s, "{} case 5u: switch(sh){{case 0u:flags[0]=1u;flags[3]=1u;break;case 1u:flags[1]=1u;flags[2]=1u;break;case 2u:flags[1]=1u;flags[3]=1u;break;case 3u:flags[2]=1u;flags[3]=1u;break;}} break;", pad).unwrap(); + writeln!(s, "{} }}", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty((unsigned int)flags[{}u]);", + pad, apc_col, flag_index + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + EmitterColumn::LoadStoreWriteData { + apc_col, + opcode_off, + shift_off, + read_data_off, + prev_data_off, + limb_index, + } => { + writeln!(s, "{}{{", pad).unwrap(); + writeln!(s, "{} unsigned char op = rec{}[{}u];", pad, i, opcode_off).unwrap(); + writeln!(s, "{} int sh = (int)rec{}[{}u];", pad, i, shift_off).unwrap(); + writeln!(s, "{} const unsigned char* rd = rec{} + {}u;", pad, i, read_data_off).unwrap(); + writeln!(s, "{} int idx = {};", pad, limb_index).unwrap(); + writeln!( + s, + "{} unsigned int prev = *(const unsigned int*)(rec{} + {}u + idx*4);", + pad, i, prev_data_off + ).unwrap(); + writeln!(s, "{} unsigned int wd = 0u;", pad).unwrap(); + writeln!(s, "{} switch (op) {{", pad).unwrap(); + writeln!(s, "{} case 0u: case 3u: wd = (unsigned int)rd[idx]; break;", pad).unwrap(); + writeln!(s, "{} case 2u: wd = (idx < 2) ? (unsigned int)rd[idx + sh] : 0u; break;", pad).unwrap(); + writeln!(s, "{} case 1u: wd = (idx == 0) ? (unsigned int)rd[sh] : 0u; break;", pad).unwrap(); + writeln!(s, "{} case 4u: wd = (idx >= sh && idx < 2 + sh) ? (unsigned int)rd[idx - sh] : prev; break;", pad).unwrap(); + writeln!(s, "{} case 5u: wd = (idx == sh) ? (unsigned int)rd[0] : prev; break;", pad).unwrap(); + writeln!(s, "{} }}", pad).unwrap(); + writeln!( + s, + "{} d_output[(unsigned long long){}u * H + r] = to_monty(wd);", + pad, apc_col + ).unwrap(); + writeln!(s, "{}}}", pad).unwrap(); + } + // Conditional handled at the wrapping level in emit_column. + EmitterColumn::Conditional { .. } => unreachable!("Conditional handled in emit_column"), + } +} + +/// Compute the BabyBear Montgomery encoding host-side. Mirrors the kernel's +/// `to_monty` exactly so verifier and unit tests compare cleanly. +pub(crate) fn host_to_monty(v: u32) -> u32 { + const M: u32 = 0x8800_0001; + const P: u32 = 0x7800_0001; + const R2: u32 = 1_172_168_163; + let x = (v as u64) * (R2 as u64); + let t = x.wrapping_mul(M as u64) & 0xFFFF_FFFF; + let u = t.wrapping_mul(P as u64); + let (x_sub_u, overflow) = x.overflowing_sub(u); + let hi = (x_sub_u >> 32) as u32; + hi + if overflow { P } else { 0 } +} + +const PRELUDE: &str = r#"// Auto-generated by powdr nvrtc_emit. Do not edit. + +__device__ __forceinline__ unsigned int monty_reduce(unsigned long long x) { + constexpr unsigned int M = 0x88000001u; + constexpr unsigned int P = 0x78000001u; + unsigned long long t = (x * (unsigned long long)M) & 0xFFFFFFFFull; + unsigned long long u = t * (unsigned long long)P; + unsigned long long x_sub_u = x - u; + bool overflow = x < u; + unsigned int hi = (unsigned int)(x_sub_u >> 32); + return hi + (overflow ? P : 0u); +} + +__device__ __forceinline__ unsigned int to_monty(unsigned int v) { + constexpr unsigned int R2 = 1172168163u; + return monty_reduce((unsigned long long)v * (unsigned long long)R2); +} + +// Helpers ported from openvm/cuda/src/apc_jit_tracegen.cu so per-column emit +// stays small. With constant offsets at every callsite, ptxas inlines and +// folds these aggressively. + +__device__ __forceinline__ void powdr_run_alu( + unsigned char opcode, + const unsigned char* b, + const unsigned char* c, + unsigned char* a +) { + if (opcode == 0u) { // ADD + unsigned int carry = 0u; + for (int i = 0; i < 4; i++) { + unsigned int s = (unsigned int)b[i] + (unsigned int)c[i] + carry; + a[i] = (unsigned char)(s & 0xFFu); + carry = s >> 8; + } + } else if (opcode == 1u) { // SUB + unsigned int borrow = 0u; + for (int i = 0; i < 4; i++) { + unsigned int rhs = (unsigned int)c[i] + borrow; + if ((unsigned int)b[i] >= rhs) { + a[i] = (unsigned char)(b[i] - rhs); + borrow = 0u; + } else { + a[i] = (unsigned char)(256u + b[i] - rhs); + borrow = 1u; + } + } + } else if (opcode == 2u) { + for (int i = 0; i < 4; i++) a[i] = b[i] ^ c[i]; + } else if (opcode == 3u) { + for (int i = 0; i < 4; i++) a[i] = b[i] | c[i]; + } else if (opcode == 4u) { + for (int i = 0; i < 4; i++) a[i] = b[i] & c[i]; + } else { + for (int i = 0; i < 4; i++) a[i] = 0u; + } +} + +__device__ __forceinline__ void powdr_get_shift_amounts( + unsigned char c0, int* limb_shift, int* bit_shift +) { + int shift = c0 % 32; + *limb_shift = shift / 8; + *bit_shift = shift % 8; +} + +// Montgomery field inverse. Input/output are in canonical form (< P). Uses +// Fermat's little theorem: x^(P-2) mod P. Hot enough only in BranchEqualDiffInv. +__device__ __forceinline__ unsigned int powdr_canonical_inv(unsigned int x) { + constexpr unsigned int P = 0x78000001u; + if (x == 0u) return 0u; + // P - 2 = 0x77FFFFFF. Square-and-multiply, MSB to LSB. + unsigned int e = 0x77FFFFFFu; + unsigned long long acc = 1ull; + unsigned long long base = (unsigned long long)x; + while (e != 0u) { + if (e & 1u) acc = (acc * base) % (unsigned long long)P; + base = (base * base) % (unsigned long long)P; + e >>= 1; + } + return (unsigned int)acc; +} + +// to_monty_canonical: produce Monty form of a canonical (< P) value. +// Equivalent to to_monty for inputs already in [0, P). +__device__ __forceinline__ unsigned int to_monty_canonical(unsigned int v) { + return to_monty(v); +} + +__device__ __forceinline__ void powdr_run_shift( + unsigned char opcode, + const unsigned char* b, + const unsigned char* c, + unsigned char* a +) { + int ls, bs; powdr_get_shift_amounts(c[0], &ls, &bs); + if (opcode == 0u) { // SLL + for (int i = 0; i < 4; i++) a[i] = 0u; + for (int i = ls; i < 4; i++) { + unsigned short high = (unsigned short)b[i - ls] << bs; + unsigned short low = (i > ls) ? ((unsigned short)b[i - ls - 1] >> (8 - bs)) : 0; + a[i] = (unsigned char)((high | low) & 0xFFu); + } + } else { // SRL / SRA + unsigned char msb = b[3] >> 7; + unsigned char fill = (opcode == 1u) ? 0u : (unsigned char)(0xFFu * msb); + for (int i = 0; i < 4; i++) a[i] = fill; + for (int i = 0; i < 4 - ls; i++) { + unsigned short p1 = (unsigned short)(b[i + ls] >> bs); + unsigned short p2v = (i + ls + 1 < 4) ? b[i + ls + 1] : fill; + unsigned short p2 = (unsigned short)p2v << (8 - bs); + a[i] = (unsigned char)((p1 | p2) & 0xFFu); + } + } +} + +"#; + +#[cfg(test)] +mod tests { + use super::*; + + fn sample_input() -> EmitterInput { + EmitterInput { + instructions: vec![EmitterInstruction { + arena_offset: 0, + record_stride: 64, + record_offset: 0, + columns: vec![ + EmitterColumn::DirectU32 { apc_col: 5, off: 16 }, + EmitterColumn::DirectU32 { apc_col: 6, off: 20 }, + ], + }], + } + } + + #[test] + fn emit_is_deterministic() { + let a = emit_jit_kernel_source(&sample_input()); + let b = emit_jit_kernel_source(&sample_input()); + assert_eq!(a.source, b.source); + assert_eq!(a.name, b.name); + assert_eq!(a.source_hash, b.source_hash); + } + + #[test] + fn emit_contains_expected_pieces() { + let k = emit_jit_kernel_source(&sample_input()); + assert!(k.source.contains("to_monty")); + assert!(k.source.contains(&k.name)); + assert!(k.source.contains("rec0 + 16u")); + assert!(k.source.contains("rec0 + 20u")); + } + + #[test] + fn fusion_groups_alu_limbs_into_single_run_alu_call() { + let input = EmitterInput { + instructions: vec![EmitterInstruction { + arena_offset: 0, + record_stride: 64, + record_offset: 0, + columns: vec![ + EmitterColumn::AluResult { apc_col: 10, opcode_off: 0, b_off: 16, c_off: 20, limb_index: 0 }, + EmitterColumn::AluResult { apc_col: 11, opcode_off: 0, b_off: 16, c_off: 20, limb_index: 1 }, + EmitterColumn::AluResult { apc_col: 12, opcode_off: 0, b_off: 16, c_off: 20, limb_index: 2 }, + EmitterColumn::AluResult { apc_col: 13, opcode_off: 0, b_off: 16, c_off: 20, limb_index: 3 }, + EmitterColumn::DirectU32 { apc_col: 5, off: 24 }, + ], + }], + }; + let k = emit_jit_kernel_source(&input); + // Find the kernel body and count run_alu invocations there (excluding + // the prelude definition). + let body = k.source.split("__global__").nth(1).expect("kernel body"); + let n_alu_calls = body.matches("powdr_run_alu(rec").count(); + assert_eq!(n_alu_calls, 1, "expected fused single ALU call, got {}\n--- source ---\n{}", n_alu_calls, k.source); + // Four stores following the call (one per limb), all from the same a[]. + let n_stores_after = body.matches("to_monty((unsigned int)a[").count(); + assert_eq!(n_stores_after, 4); + } + + #[test] + fn host_to_monty_matches_known_values() { + assert_eq!(host_to_monty(0), 0); + assert_eq!(host_to_monty(1), 268_435_454); + let p: u32 = 0x7800_0001; + let r_mod_p: u32 = 268_435_454; + assert_eq!(host_to_monty(p - 1), p - r_mod_p); + } +} diff --git a/openvm/src/powdr_extension/trace_generator/cuda/periphery.rs b/openvm/src/powdr_extension/trace_generator/cuda/periphery.rs index 33553ec958..5daebe38c9 100644 --- a/openvm/src/powdr_extension/trace_generator/cuda/periphery.rs +++ b/openvm/src/powdr_extension/trace_generator/cuda/periphery.rs @@ -10,9 +10,9 @@ use openvm_circuit_primitives::{ range_tuple::{RangeTupleCheckerAir, RangeTupleCheckerChip, RangeTupleCheckerChipGPU}, var_range::{VariableRangeCheckerAir, VariableRangeCheckerChipGPU}, }; -use openvm_cuda_backend::engine::GpuBabyBearPoseidon2Engine; -use openvm_cuda_backend::prover_backend::GpuBackend; -use openvm_stark_backend::{config::StarkGenericConfig, p3_field::PrimeField32}; +use openvm_cuda_backend::BabyBearPoseidon2GpuEngine as GpuBabyBearPoseidon2CpuEngine; +use openvm_cuda_backend::GpuBackend; +use openvm_stark_backend::{p3_field::PrimeField32, StarkProtocolConfig}; use crate::{ isa::OpenVmISA, powdr_extension::trace_generator::common::DummyExecutor, BabyBearSC, @@ -84,7 +84,9 @@ impl PowdrPeripheryInstancesGpu { } } -impl VmExecutionExtension for SharedPeripheryChipsGpu { +impl, ISA: OpenVmISA> + VmExecutionExtension for SharedPeripheryChipsGpu +{ type Executor = DummyExecutor; fn extend_execution( @@ -96,7 +98,7 @@ impl VmExecutionExtension for SharedPeripher } } -impl VmCircuitExtension +impl VmCircuitExtension for SharedPeripheryChipsGpu { fn extend_circuit(&self, inventory: &mut AirInventory) -> Result<(), AirInventoryError> { @@ -148,7 +150,7 @@ impl VmCircuitExtension pub struct SharedPeripheryChipsGpuProverExt; impl - VmProverExtension> + VmProverExtension> for SharedPeripheryChipsGpuProverExt { fn extend_prover( diff --git a/openvm/src/powdr_extension/trace_generator/jit_mapping.rs b/openvm/src/powdr_extension/trace_generator/jit_mapping.rs new file mode 100644 index 0000000000..c5bf441a65 --- /dev/null +++ b/openvm/src/powdr_extension/trace_generator/jit_mapping.rs @@ -0,0 +1,1908 @@ +//! Column-to-record mapping tables for JIT trace generation. +//! +//! Each instruction type (AIR) has a mapping from trace column indices to +//! computations on the raw record bytes stored in the MatrixRecordArena. +//! This allows us to compute only the surviving APC columns directly from +//! record data, bypassing the full `fill_trace_row` pipeline. + +use openvm_stark_backend::p3_field::PrimeField32; + +/// How to compute a single trace column value from raw record bytes. +#[derive(Debug, Clone)] +pub enum ColumnComputation { + /// Read a u32 field at the given byte offset in the record. + /// Produces `F::from_u32(*(u32*)(record + offset))`. + DirectU32 { record_byte_offset: usize }, + + /// Read a single byte at the given byte offset in the record. + /// Produces `F::from_u8(record[offset])`. + DirectU8 { record_byte_offset: usize }, + + /// Timestamp decomposition: computes one limb of + /// `decompose((record[curr_ts_offset] + delta) - record[prev_ts_offset] - 1, max_bits)`. + /// `max_bits` comes from the VariableRangeChecker configuration at runtime. + /// `limb_index` selects which limb (0 or 1). + /// `delta` accounts for the timestamp increment within the instruction + /// (e.g., 0 for first read, 1 for second read, 2 for write). + TimestampDecomp { + /// Byte offset of the "current" timestamp (u32) in the record. + curr_ts_byte_offset: usize, + /// Delta added to the current timestamp before decomposing. + curr_ts_delta: u32, + /// Byte offset of the "previous" timestamp (u32) in the record. + prev_ts_byte_offset: usize, + /// Which limb of the decomposition (0 = low, 1 = high). + limb_index: usize, + }, + + /// ALU result: computes `run_alu(opcode, b, c)[limb_index]`. + /// `opcode_byte_offset` points to the local_opcode u8 field. + /// `b_byte_offset` and `c_byte_offset` point to the start of the b[4] and c[4] arrays. + AluResult { + opcode_byte_offset: usize, + b_byte_offset: usize, + c_byte_offset: usize, + limb_index: usize, + }, + + /// Boolean from opcode: `F::from_bool(record[opcode_byte_offset] == expected_opcode)`. + BoolFromOpcode { + opcode_byte_offset: usize, + expected_opcode: u8, + }, + + /// Conditional: if `record[condition_byte_offset] != 0`, use `then_comp`; else `F::ZERO`. + Conditional { + condition_byte_offset: usize, + then_comp: Box, + }, + + /// Read a u16 field at the given byte offset in the record. + DirectU16 { record_byte_offset: usize }, + + /// Pointer limb: computes `uint16_t limb of (val + sign_extend(imm, imm_sign))`. + PointerLimb { + val_byte_offset: usize, + imm_byte_offset: usize, + imm_sign_byte_offset: usize, + limb_index: usize, + }, + + /// LoadStore: rd_rs2_ptr (0 if UINT32_MAX, otherwise the value) + LoadStoreRdRs2Ptr { rd_rs2_ptr_byte_offset: usize }, + + /// LoadStore: needs_write (rd_rs2_ptr != UINT32_MAX) + LoadStoreNeedsWrite { rd_rs2_ptr_byte_offset: usize }, + + /// LoadStore: write_base_aux.prev_timestamp (conditional on needs_write) + LoadStoreWriteAuxPrevTs { + write_prev_ts_byte_offset: usize, + rd_rs2_ptr_byte_offset: usize, + }, + + /// LoadStore: write_base_aux timestamp decomposition (conditional on needs_write) + LoadStoreWriteAuxDecomp { + from_ts_byte_offset: usize, + write_prev_ts_byte_offset: usize, + rd_rs2_ptr_byte_offset: usize, + limb_index: usize, + }, + + /// LoadStore: is_load flag + LoadStoreIsLoad { opcode_byte_offset: usize }, + + /// LoadStore: flags[4] based on opcode and shift + LoadStoreFlag { + opcode_byte_offset: usize, + shift_byte_offset: usize, + flag_index: usize, + }, + + /// LoadStore: write_data computed from opcode, shift, read_data, prev_data + LoadStoreWriteData { + opcode_byte_offset: usize, + shift_byte_offset: usize, + read_data_byte_offset: usize, + prev_data_byte_offset: usize, + limb_index: usize, + }, + + /// Shift result: a[limb_index] from shift operation + ShiftResult { opcode_byte_offset: usize, b_byte_offset: usize, c_byte_offset: usize, limb_index: usize }, + /// Shift: bit_multiplier_left = is_sll ? (1 << bit_shift) : 0 + ShiftBitMulLeft { opcode_byte_offset: usize, c_byte_offset: usize }, + /// Shift: bit_multiplier_right = !is_sll ? (1 << bit_shift) : 0 + ShiftBitMulRight { opcode_byte_offset: usize, c_byte_offset: usize }, + /// Shift: b_sign = is_sra ? (b[3] >> 7) : 0 + ShiftBSign { opcode_byte_offset: usize, b_byte_offset: usize }, + /// Shift: bit_shift_marker[marker_index] = (bit_shift == marker_index) + ShiftBitMarker { c_byte_offset: usize, marker_index: usize }, + /// Shift: limb_shift_marker[marker_index] = (limb_shift == marker_index) + ShiftLimbMarker { c_byte_offset: usize, marker_index: usize }, + /// Shift: bit_shift_carry[limb_index] + ShiftBitCarry { opcode_byte_offset: usize, b_byte_offset: usize, c_byte_offset: usize, limb_index: usize }, + + /// BranchEqual: cmp_result + BranchEqualCmpResult { a_byte_offset: usize, b_byte_offset: usize, opcode_byte_offset: usize }, + /// BranchEqual: diff_inv_marker[marker_index] + BranchEqualDiffInvMarker { a_byte_offset: usize, b_byte_offset: usize, opcode_byte_offset: usize, marker_index: usize }, + + // ── LessThanCoreAir<4, 8> (slt / sltu / slti / sltiu) ──────────────────── + // All five arms share the (opcode, b, c) byte-offset signature so a future + // NVRTC fusion pass can recognise them as a single group and emit one + // run_less_than per row, indexing into its result. + /// LessThan: cmp_result (0 or 1). + LessThanCmpResult { opcode_byte_offset: usize, b_byte_offset: usize, c_byte_offset: usize }, + /// LessThan: diff_val (canonical-form field value). + LessThanDiffVal { opcode_byte_offset: usize, b_byte_offset: usize, c_byte_offset: usize }, + /// LessThan: diff_marker[marker_index]. + LessThanDiffMarker { opcode_byte_offset: usize, b_byte_offset: usize, c_byte_offset: usize, marker_index: usize }, + /// LessThan: b_msb_f (signed-aware MSB encoding of b[3]). + LessThanBMsbF { opcode_byte_offset: usize, b_byte_offset: usize }, + /// LessThan: c_msb_f. + LessThanCMsbF { opcode_byte_offset: usize, c_byte_offset: usize }, + + // ── Rv32AuipcCoreAir (auipc) — Rv32RdWriteAdapter ─────────────────────── + /// AUIPC rd_data limb: byte `limb_index` of `(pc + (imm << 8))`. + /// pc and imm are stored as u32 in the record. + AuipcRdLimb { pc_byte_offset: usize, imm_byte_offset: usize, limb_index: usize }, + + // ── Rv32JalrCoreAir (jalr) — Rv32JalrAdapter ──────────────────────────── + /// JALR `to_pc_least_sig_bit`: `(rs1 + imm + (imm_sign ? 0xFFFF_0000 : 0)) & 1`. + JalrToPcLsb { + rs1_byte_offset: usize, + imm_byte_offset: usize, + imm_sign_byte_offset: usize, + }, + /// JALR `to_pc_limbs[limb_index]`. Layout: + /// - limb 0: `(to_pc & 0xFFFF) >> 1` + /// - limb 1: `to_pc >> 16` + JalrToPcLimb { + rs1_byte_offset: usize, + imm_byte_offset: usize, + imm_sign_byte_offset: usize, + limb_index: usize, + }, + /// JALR rd_data limb: byte `limb_index + 1` of `(pc + 4)` for the top 3 + /// limbs the chip writes (`limb_index ∈ 0..3`). + JalrRdLimb { pc_byte_offset: usize, limb_index: usize }, + + /// Like [`Conditional`] but the gate is "u32 at `ptr_byte_offset` + /// is not `u32::MAX`" — used for adapter rd_aux fields that are only + /// populated when the chip actually performs a write. + ConditionalNotMaxU32 { + ptr_byte_offset: usize, + then_comp: Box, + }, + + // ── BranchLessThanCoreAir<4, 8> (blt / bltu / bge / bgeu) ──────────────── + /// BranchLt: cmp_result (1 if branch taken). + BranchLtCmpResult { opcode_byte_offset: usize, a_byte_offset: usize, b_byte_offset: usize }, + /// BranchLt: cmp_lt (1 if a < b irrespective of bge/blt). + BranchLtCmpLt { opcode_byte_offset: usize, a_byte_offset: usize, b_byte_offset: usize }, + /// BranchLt: diff_val. + BranchLtDiffVal { opcode_byte_offset: usize, a_byte_offset: usize, b_byte_offset: usize }, + /// BranchLt: diff_marker[marker_index]. + BranchLtDiffMarker { opcode_byte_offset: usize, a_byte_offset: usize, b_byte_offset: usize, marker_index: usize }, + /// BranchLt: a_msb_f. + BranchLtAMsbF { opcode_byte_offset: usize, a_byte_offset: usize }, + /// BranchLt: b_msb_f. + BranchLtBMsbF { opcode_byte_offset: usize, b_byte_offset: usize }, + + /// Constant value. + Constant(u32), +} + +/// A single column mapping entry: which column index maps to which computation. +#[derive(Debug, Clone)] +pub struct ColumnMapping { + /// The column index in the original AIR trace (= `original_poly_index` in Substitution). + pub col_index: usize, + /// How to compute this column from the record bytes. + pub computation: ColumnComputation, +} + +/// Complete mapping table for one AIR type. +#[derive(Debug, Clone)] +pub struct AirColumnMapping { + /// Human-readable AIR name for debugging. + pub air_name: &'static str, + /// Total width (number of columns) of this AIR. + pub width: usize, + /// Size of the record in bytes (adapter + core, including padding). + pub record_byte_size: usize, + /// Mapping for each column. Indexed by column index. + pub columns: Vec, +} + +/// Evaluate a column computation given the raw record bytes. +pub fn eval_column(comp: &ColumnComputation, record: &[u8], range_max_bits: u32) -> F { + match comp { + ColumnComputation::DirectU32 { record_byte_offset } => { + let bytes = &record[*record_byte_offset..*record_byte_offset + 4]; + let val = u32::from_le_bytes(bytes.try_into().unwrap()); + F::from_u32(val) + } + ColumnComputation::DirectU8 { record_byte_offset } => { + F::from_u8(record[*record_byte_offset]) + } + ColumnComputation::TimestampDecomp { + curr_ts_byte_offset, + curr_ts_delta, + prev_ts_byte_offset, + limb_index, + } => { + let curr = u32::from_le_bytes( + record[*curr_ts_byte_offset..*curr_ts_byte_offset + 4] + .try_into() + .unwrap(), + ); + let prev = u32::from_le_bytes( + record[*prev_ts_byte_offset..*prev_ts_byte_offset + 4] + .try_into() + .unwrap(), + ); + let diff = (curr + curr_ts_delta).wrapping_sub(prev).wrapping_sub(1); + let mask = (1u32 << range_max_bits) - 1; + let limb = (diff >> (range_max_bits * *limb_index as u32)) & mask; + F::from_u32(limb) + } + ColumnComputation::AluResult { + opcode_byte_offset, + b_byte_offset, + c_byte_offset, + limb_index, + } => { + let opcode = record[*opcode_byte_offset]; + let b = &record[*b_byte_offset..*b_byte_offset + 4]; + let c = &record[*c_byte_offset..*c_byte_offset + 4]; + let result = run_alu_byte(opcode, b, c); + F::from_u8(result[*limb_index]) + } + ColumnComputation::BoolFromOpcode { + opcode_byte_offset, + expected_opcode, + } => F::from_bool(record[*opcode_byte_offset] == *expected_opcode), + ColumnComputation::Conditional { + condition_byte_offset, + then_comp, + } => { + if record[*condition_byte_offset] != 0 { + eval_column(then_comp, record, range_max_bits) + } else { + F::ZERO + } + } + ColumnComputation::DirectU16 { record_byte_offset } => { + let bytes = &record[*record_byte_offset..*record_byte_offset + 2]; + let val = u16::from_le_bytes(bytes.try_into().unwrap()); + F::from_u32(val as u32) + } + ColumnComputation::PointerLimb { + val_byte_offset, + imm_byte_offset, + imm_sign_byte_offset, + limb_index, + } => { + let val = u32::from_le_bytes( + record[*val_byte_offset..*val_byte_offset + 4].try_into().unwrap(), + ); + let imm = u16::from_le_bytes( + record[*imm_byte_offset..*imm_byte_offset + 2].try_into().unwrap(), + ) as u32; + let sign = record[*imm_sign_byte_offset]; + let imm_ext = imm + if sign != 0 { 0xFFFF0000u32 } else { 0 }; + let ptr = val.wrapping_add(imm_ext); + let limb = ((ptr >> (16 * *limb_index as u32)) & 0xFFFF) as u32; + F::from_u32(limb) + } + ColumnComputation::LoadStoreRdRs2Ptr { rd_rs2_ptr_byte_offset } => { + let val = u32::from_le_bytes( + record[*rd_rs2_ptr_byte_offset..*rd_rs2_ptr_byte_offset + 4] + .try_into().unwrap(), + ); + if val == u32::MAX { F::ZERO } else { F::from_u32(val) } + } + ColumnComputation::LoadStoreNeedsWrite { rd_rs2_ptr_byte_offset } => { + let val = u32::from_le_bytes( + record[*rd_rs2_ptr_byte_offset..*rd_rs2_ptr_byte_offset + 4] + .try_into().unwrap(), + ); + F::from_bool(val != u32::MAX) + } + ColumnComputation::LoadStoreWriteAuxPrevTs { + write_prev_ts_byte_offset, + rd_rs2_ptr_byte_offset, + } => { + let rd_rs2 = u32::from_le_bytes( + record[*rd_rs2_ptr_byte_offset..*rd_rs2_ptr_byte_offset + 4] + .try_into().unwrap(), + ); + if rd_rs2 == u32::MAX { + F::ZERO + } else { + let val = u32::from_le_bytes( + record[*write_prev_ts_byte_offset..*write_prev_ts_byte_offset + 4] + .try_into().unwrap(), + ); + F::from_u32(val) + } + } + ColumnComputation::LoadStoreWriteAuxDecomp { + from_ts_byte_offset, + write_prev_ts_byte_offset, + rd_rs2_ptr_byte_offset, + limb_index, + } => { + let rd_rs2 = u32::from_le_bytes( + record[*rd_rs2_ptr_byte_offset..*rd_rs2_ptr_byte_offset + 4] + .try_into().unwrap(), + ); + if rd_rs2 == u32::MAX { + F::ZERO + } else { + let curr = u32::from_le_bytes( + record[*from_ts_byte_offset..*from_ts_byte_offset + 4] + .try_into().unwrap(), + ); + let prev = u32::from_le_bytes( + record[*write_prev_ts_byte_offset..*write_prev_ts_byte_offset + 4] + .try_into().unwrap(), + ); + let diff = (curr + 2).wrapping_sub(prev).wrapping_sub(1); + let mask = (1u32 << range_max_bits) - 1; + let limb = (diff >> (range_max_bits * *limb_index as u32)) & mask; + F::from_u32(limb) + } + } + ColumnComputation::LoadStoreIsLoad { opcode_byte_offset } => { + let opcode = record[*opcode_byte_offset]; + // LOADW=0, LOADBU=1, LOADHU=2 are loads + F::from_bool(opcode <= 2) + } + ColumnComputation::LoadStoreFlag { + opcode_byte_offset, + shift_byte_offset, + flag_index, + } => { + let opcode = record[*opcode_byte_offset]; + let shift = record[*shift_byte_offset]; + let flags = compute_loadstore_flags(opcode, shift); + F::from_u32(flags[*flag_index] as u32) + } + ColumnComputation::LoadStoreWriteData { + opcode_byte_offset, + shift_byte_offset, + read_data_byte_offset, + prev_data_byte_offset, + limb_index, + } => { + let opcode = record[*opcode_byte_offset]; + let shift = record[*shift_byte_offset] as usize; + let rd = &record[*read_data_byte_offset..*read_data_byte_offset + 4]; + let prev = [ + u32::from_le_bytes(record[*prev_data_byte_offset..*prev_data_byte_offset + 4].try_into().unwrap()), + u32::from_le_bytes(record[*prev_data_byte_offset + 4..*prev_data_byte_offset + 8].try_into().unwrap()), + u32::from_le_bytes(record[*prev_data_byte_offset + 8..*prev_data_byte_offset + 12].try_into().unwrap()), + u32::from_le_bytes(record[*prev_data_byte_offset + 12..*prev_data_byte_offset + 16].try_into().unwrap()), + ]; + let wd = compute_loadstore_write_data(opcode, shift, rd, &prev); + F::from_u32(wd[*limb_index]) + } + ColumnComputation::ShiftResult { opcode_byte_offset, b_byte_offset, c_byte_offset, limb_index } => { + let opcode = record[*opcode_byte_offset]; + let b: [u8; 4] = record[*b_byte_offset..*b_byte_offset + 4].try_into().unwrap(); + let c: [u8; 4] = record[*c_byte_offset..*c_byte_offset + 4].try_into().unwrap(); + let a = run_shift(opcode, &b, &c); + F::from_u8(a[*limb_index]) + } + ColumnComputation::ShiftBitMulLeft { opcode_byte_offset, c_byte_offset } => { + let is_sll = record[*opcode_byte_offset] == 0; + let (_, bit_shift) = get_shift_amounts(record[*c_byte_offset]); + if is_sll { F::from_u32(1u32 << bit_shift) } else { F::ZERO } + } + ColumnComputation::ShiftBitMulRight { opcode_byte_offset, c_byte_offset } => { + let is_sll = record[*opcode_byte_offset] == 0; + let (_, bit_shift) = get_shift_amounts(record[*c_byte_offset]); + if !is_sll { F::from_u32(1u32 << bit_shift) } else { F::ZERO } + } + ColumnComputation::ShiftBSign { opcode_byte_offset, b_byte_offset } => { + let is_sra = record[*opcode_byte_offset] == 2; + if is_sra { F::from_u32((record[*b_byte_offset + 3] >> 7) as u32) } else { F::ZERO } + } + ColumnComputation::ShiftBitMarker { c_byte_offset, marker_index } => { + let (_, bit_shift) = get_shift_amounts(record[*c_byte_offset]); + F::from_bool(bit_shift == *marker_index) + } + ColumnComputation::ShiftLimbMarker { c_byte_offset, marker_index } => { + let (limb_shift, _) = get_shift_amounts(record[*c_byte_offset]); + F::from_bool(limb_shift == *marker_index) + } + ColumnComputation::ShiftBitCarry { opcode_byte_offset, b_byte_offset, c_byte_offset, limb_index } => { + let is_sll = record[*opcode_byte_offset] == 0; + let (_, bit_shift) = get_shift_amounts(record[*c_byte_offset]); + if bit_shift == 0 { + F::ZERO + } else { + let b_val = record[*b_byte_offset + *limb_index]; + let carry = if is_sll { + b_val >> (8 - bit_shift as u8) + } else { + b_val & ((1u8 << bit_shift as u8) - 1) + }; + F::from_u8(carry) + } + } + ColumnComputation::BranchEqualCmpResult { a_byte_offset, b_byte_offset, opcode_byte_offset } => { + let a = &record[*a_byte_offset..*a_byte_offset + 4]; + let b = &record[*b_byte_offset..*b_byte_offset + 4]; + let is_beq = record[*opcode_byte_offset] == 0; + let are_equal = a == b; + F::from_bool(if is_beq { are_equal } else { !are_equal }) + } + ColumnComputation::BranchEqualDiffInvMarker { a_byte_offset, b_byte_offset, opcode_byte_offset, marker_index } => { + let a = &record[*a_byte_offset..*a_byte_offset + 4]; + let b = &record[*b_byte_offset..*b_byte_offset + 4]; + // Find first differing limb + let mut diff_idx = 4usize; // = NUM_LIMBS means equal + for i in 0..4 { + if a[i] != b[i] { diff_idx = i; break; } + } + if diff_idx == 4 { diff_idx = 0; } // when equal, marker at index 0 + if *marker_index == diff_idx && a != b { + // diff_inv = inv(a[diff_idx] - b[diff_idx]) in the field + let a_val = F::from_u8(a[diff_idx]); + let b_val = F::from_u8(b[diff_idx]); + (a_val - b_val).inverse() + } else { + F::ZERO + } + } + ColumnComputation::LessThanCmpResult { opcode_byte_offset, b_byte_offset, c_byte_offset } => { + let is_slt = record[*opcode_byte_offset] == 0; + let b: &[u8] = &record[*b_byte_offset..*b_byte_offset + 4]; + let c: &[u8] = &record[*c_byte_offset..*c_byte_offset + 4]; + let (cmp_result, _, _, _) = run_less_than_host(is_slt, b, c); + F::from_bool(cmp_result) + } + ColumnComputation::LessThanDiffVal { opcode_byte_offset, b_byte_offset, c_byte_offset } => { + let is_slt = record[*opcode_byte_offset] == 0; + let b: &[u8] = &record[*b_byte_offset..*b_byte_offset + 4]; + let c: &[u8] = &record[*c_byte_offset..*c_byte_offset + 4]; + let (cmp_result, diff_idx, b_sign, c_sign) = run_less_than_host(is_slt, b, c); + less_than_diff_val::(cmp_result, diff_idx, b, c, b_sign, c_sign) + } + ColumnComputation::LessThanDiffMarker { opcode_byte_offset, b_byte_offset, c_byte_offset, marker_index } => { + let is_slt = record[*opcode_byte_offset] == 0; + let b: &[u8] = &record[*b_byte_offset..*b_byte_offset + 4]; + let c: &[u8] = &record[*c_byte_offset..*c_byte_offset + 4]; + let (_, diff_idx, _, _) = run_less_than_host(is_slt, b, c); + // marker[diff_idx] = 1 only when a != b (diff_idx != 4); else all zero. + if diff_idx != 4 && diff_idx == *marker_index { F::ONE } else { F::ZERO } + } + ColumnComputation::LessThanBMsbF { opcode_byte_offset, b_byte_offset } => { + let is_slt = record[*opcode_byte_offset] == 0; + let b3 = record[*b_byte_offset + 3]; + // b_sign = signed-mode AND high bit set. + let b_sign = is_slt && (b3 >> 7) == 1; + if b_sign { + // -F::from_u16((1 << 8) - b3) + F::ZERO - F::from_u16(256 - b3 as u16) + } else { + F::from_u8(b3) + } + } + ColumnComputation::LessThanCMsbF { opcode_byte_offset, c_byte_offset } => { + let is_slt = record[*opcode_byte_offset] == 0; + let c3 = record[*c_byte_offset + 3]; + let c_sign = is_slt && (c3 >> 7) == 1; + if c_sign { + F::ZERO - F::from_u16(256 - c3 as u16) + } else { + F::from_u8(c3) + } + } + ColumnComputation::AuipcRdLimb { pc_byte_offset, imm_byte_offset, limb_index } => { + let pc = read_u32(record, *pc_byte_offset); + let imm = read_u32(record, *imm_byte_offset); + let rd = pc.wrapping_add(imm << 8); + F::from_u8(rd.to_le_bytes()[*limb_index]) + } + ColumnComputation::JalrToPcLsb { rs1_byte_offset, imm_byte_offset, imm_sign_byte_offset } => { + let rs1 = read_u32(record, *rs1_byte_offset); + let imm = read_u16(record, *imm_byte_offset) as u32; + let sign = record[*imm_sign_byte_offset] != 0; + let to_pc = rs1.wrapping_add(imm + if sign { 0xFFFF_0000 } else { 0 }); + F::from_u8((to_pc & 1) as u8) + } + ColumnComputation::JalrToPcLimb { rs1_byte_offset, imm_byte_offset, imm_sign_byte_offset, limb_index } => { + let rs1 = read_u32(record, *rs1_byte_offset); + let imm = read_u16(record, *imm_byte_offset) as u32; + let sign = record[*imm_sign_byte_offset] != 0; + let to_pc = rs1.wrapping_add(imm + if sign { 0xFFFF_0000 } else { 0 }); + let v = match *limb_index { + 0 => (to_pc & 0xFFFF) >> 1, + 1 => to_pc >> 16, + _ => panic!("JalrToPcLimb: invalid limb_index {limb_index}"), + }; + F::from_u32(v) + } + ColumnComputation::JalrRdLimb { pc_byte_offset, limb_index } => { + let pc = read_u32(record, *pc_byte_offset); + // Chip stores top 3 limbs of (pc + 4) — limb_index ∈ 0..3 + // mapping to byte (limb_index + 1) of the little-endian encoding. + let rd = pc.wrapping_add(4); + F::from_u8(rd.to_le_bytes()[*limb_index + 1]) + } + ColumnComputation::ConditionalNotMaxU32 { ptr_byte_offset, then_comp } => { + let v = read_u32(record, *ptr_byte_offset); + if v == u32::MAX { + F::ZERO + } else { + eval_column(then_comp, record, range_max_bits) + } + } + ColumnComputation::BranchLtCmpResult { opcode_byte_offset, a_byte_offset, b_byte_offset } => { + let op = record[*opcode_byte_offset]; + let a: &[u8] = &record[*a_byte_offset..*a_byte_offset + 4]; + let b: &[u8] = &record[*b_byte_offset..*b_byte_offset + 4]; + let (cmp_result, _, _, _) = run_branch_lt_host(op, a, b); + F::from_bool(cmp_result) + } + ColumnComputation::BranchLtCmpLt { opcode_byte_offset, a_byte_offset, b_byte_offset } => { + let op = record[*opcode_byte_offset]; + let a: &[u8] = &record[*a_byte_offset..*a_byte_offset + 4]; + let b: &[u8] = &record[*b_byte_offset..*b_byte_offset + 4]; + let (cmp_result, _, _, _) = run_branch_lt_host(op, a, b); + // ge_op: BGE=2 or BGEU=3 + let ge_op = op == 2 || op == 3; + F::from_bool(cmp_result ^ ge_op) + } + ColumnComputation::BranchLtDiffVal { opcode_byte_offset, a_byte_offset, b_byte_offset } => { + let op = record[*opcode_byte_offset]; + let a: &[u8] = &record[*a_byte_offset..*a_byte_offset + 4]; + let b: &[u8] = &record[*b_byte_offset..*b_byte_offset + 4]; + let (cmp_result, diff_idx, a_sign, b_sign) = run_branch_lt_host(op, a, b); + let ge_op = op == 2 || op == 3; + let cmp_lt = cmp_result ^ ge_op; + branch_lt_diff_val::(cmp_lt, diff_idx, a, b, a_sign, b_sign) + } + ColumnComputation::BranchLtDiffMarker { opcode_byte_offset, a_byte_offset, b_byte_offset, marker_index } => { + let op = record[*opcode_byte_offset]; + let a: &[u8] = &record[*a_byte_offset..*a_byte_offset + 4]; + let b: &[u8] = &record[*b_byte_offset..*b_byte_offset + 4]; + let (_, diff_idx, _, _) = run_branch_lt_host(op, a, b); + if diff_idx != 4 && diff_idx == *marker_index { F::ONE } else { F::ZERO } + } + ColumnComputation::BranchLtAMsbF { opcode_byte_offset, a_byte_offset } => { + let op = record[*opcode_byte_offset]; + let signed = op == 0 || op == 2; // BLT or BGE + let a3 = record[*a_byte_offset + 3]; + let a_sign = signed && (a3 >> 7) == 1; + if a_sign { + F::ZERO - F::from_u16(256 - a3 as u16) + } else { + F::from_u8(a3) + } + } + ColumnComputation::BranchLtBMsbF { opcode_byte_offset, b_byte_offset } => { + let op = record[*opcode_byte_offset]; + let signed = op == 0 || op == 2; + let b3 = record[*b_byte_offset + 3]; + let b_sign = signed && (b3 >> 7) == 1; + if b_sign { + F::ZERO - F::from_u16(256 - b3 as u16) + } else { + F::from_u8(b3) + } + } + ColumnComputation::Constant(val) => F::from_u32(*val), + } +} + +/// Get shift amounts from c[0]: (limb_shift, bit_shift) +fn get_shift_amounts(c0: u8) -> (usize, usize) { + let max_bits = 4 * 8; // NUM_LIMBS * CELL_BITS = 32 + let shift = (c0 as usize) % max_bits; + (shift / 8, shift % 8) +} + +/// Run shift operation, returning result a[4] +fn run_shift(opcode: u8, b: &[u8; 4], c: &[u8; 4]) -> [u8; 4] { + let (limb_shift, bit_shift) = get_shift_amounts(c[0]); + let mut a = [0u8; 4]; + + if opcode == 0 { + // SLL + for i in limb_shift..4 { + if i > limb_shift { + let high = (b[i - limb_shift] as u16) << bit_shift as u16; + let low = (b[i - limb_shift - 1] as u16) >> (8 - bit_shift) as u16; + a[i] = ((high | low) & 0xFF) as u8; + } else { + a[i] = (((b[i - limb_shift] as u16) << bit_shift as u16) & 0xFF) as u8; + } + } + } else { + // SRL or SRA + let is_logical = opcode == 1; + let msb = b[3] >> 7; + let fill: u8 = if is_logical { 0 } else { 0xFF * msb }; + for i in 0..4 { a[i] = fill; } + let limit = 4 - limb_shift; + for i in 0..limit { + let part1 = (b[i + limb_shift] >> bit_shift as u8) as u16; + let part2_val = if i + limb_shift + 1 < 4 { b[i + limb_shift + 1] } else { fill }; + let part2 = (part2_val as u16) << (8 - bit_shift) as u16; + a[i] = ((part1 | part2) & 0xFF) as u8; + } + } + a +} + +/// Read a little-endian u32 from `record[offset..offset+4]`. +#[inline] +fn read_u32(record: &[u8], offset: usize) -> u32 { + u32::from_le_bytes(record[offset..offset + 4].try_into().unwrap()) +} + +/// Read a little-endian u16 from `record[offset..offset+2]`. +#[inline] +fn read_u16(record: &[u8], offset: usize) -> u16 { + u16::from_le_bytes(record[offset..offset + 2].try_into().unwrap()) +} + +/// Mirror of `LessThanCoreAir::run_less_than` for `NUM_LIMBS=4, LIMB_BITS=8`. +/// Returns `(cmp_result, diff_idx, b_sign, c_sign)`. `diff_idx` is the most +/// significant index where `b[i] != c[i]` (walking from MSB), or 4 when equal. +fn run_less_than_host(is_slt: bool, b: &[u8], c: &[u8]) -> (bool, usize, bool, bool) { + let b_sign = is_slt && (b[3] >> 7) == 1; + let c_sign = is_slt && (c[3] >> 7) == 1; + for i in (0..4).rev() { + if b[i] != c[i] { + return ((b[i] < c[i]) ^ b_sign ^ c_sign, i, b_sign, c_sign); + } + } + (false, 4, b_sign, c_sign) +} + +/// Mirror of `BranchLessThanCoreAir::run_cmp` for `NUM_LIMBS=4, LIMB_BITS=8`. +/// Local opcodes: BLT=0, BLTU=1, BGE=2, BGEU=3. +fn run_branch_lt_host(local_opcode: u8, a: &[u8], b: &[u8]) -> (bool, usize, bool, bool) { + let signed = local_opcode == 0 /* BLT */ || local_opcode == 2 /* BGE */; + let ge_op = local_opcode == 2 /* BGE */ || local_opcode == 3 /* BGEU */; + let a_sign = signed && (a[3] >> 7) == 1; + let b_sign = signed && (b[3] >> 7) == 1; + for i in (0..4).rev() { + if a[i] != b[i] { + return ((a[i] < b[i]) ^ a_sign ^ b_sign ^ ge_op, i, a_sign, b_sign); + } + } + (ge_op, 4, a_sign, b_sign) +} + +/// Compute LessThan's `diff_val` field-element value given the precomputed +/// `(cmp_result, diff_idx, b_sign, c_sign)` from `run_less_than_host`. +fn less_than_diff_val( + cmp_result: bool, + diff_idx: usize, + b: &[u8], + c: &[u8], + b_sign: bool, + c_sign: bool, +) -> F { + if diff_idx == 4 { + return F::ZERO; + } + if diff_idx == 3 { + // Use signed-aware MSb encodings + let b_msb = if b_sign { F::ZERO - F::from_u16(256 - b[3] as u16) } else { F::from_u8(b[3]) }; + let c_msb = if c_sign { F::ZERO - F::from_u16(256 - c[3] as u16) } else { F::from_u8(c[3]) }; + if cmp_result { c_msb - b_msb } else { b_msb - c_msb } + } else if cmp_result { + F::from_u8(c[diff_idx] - b[diff_idx]) + } else { + F::from_u8(b[diff_idx] - c[diff_idx]) + } +} + +/// Compute BranchLt's `diff_val` field-element value. +fn branch_lt_diff_val( + cmp_lt: bool, + diff_idx: usize, + a: &[u8], + b: &[u8], + a_sign: bool, + b_sign: bool, +) -> F { + if diff_idx == 4 { + return F::ZERO; + } + if diff_idx == 3 { + let a_msb = if a_sign { F::ZERO - F::from_u16(256 - a[3] as u16) } else { F::from_u8(a[3]) }; + let b_msb = if b_sign { F::ZERO - F::from_u16(256 - b[3] as u16) } else { F::from_u8(b[3]) }; + if cmp_lt { b_msb - a_msb } else { a_msb - b_msb } + } else if cmp_lt { + F::from_u8(b[diff_idx] - a[diff_idx]) + } else { + F::from_u8(a[diff_idx] - b[diff_idx]) + } +} + +/// Compute LoadStore flags[4] based on opcode and shift. +fn compute_loadstore_flags(opcode: u8, shift: u8) -> [u8; 4] { + let mut flags = [0u8; 4]; + match opcode { + 0 => flags[0] = 2, // LOADW + 2 => match shift { 0 => flags[1] = 2, 2 => flags[2] = 2, _ => {} } // LOADHU + 1 => match shift { // LOADBU + 0 => flags[3] = 2, 1 => flags[0] = 1, 2 => flags[1] = 1, 3 => flags[2] = 1, _ => {} + } + 3 => flags[3] = 1, // STOREW + 4 => match shift { // STOREH + 0 => { flags[0] = 1; flags[1] = 1; } + 2 => { flags[0] = 1; flags[2] = 1; } + _ => {} + } + 5 => match shift { // STOREB + 0 => { flags[0] = 1; flags[3] = 1; } + 1 => { flags[1] = 1; flags[2] = 1; } + 2 => { flags[1] = 1; flags[3] = 1; } + 3 => { flags[2] = 1; flags[3] = 1; } + _ => {} + } + _ => {} // LOADB=6, LOADH=7 (sign extension variants) + } + flags +} + +/// Compute LoadStore write_data[4] based on opcode, shift, read_data, prev_data. +fn compute_loadstore_write_data(opcode: u8, shift: usize, rd: &[u8], prev: &[u32; 4]) -> [u32; 4] { + let mut wd = [0u32; 4]; + match opcode { + 0 => { // LOADW + for i in 0..4 { wd[i] = rd[i] as u32; } + } + 2 => { // LOADHU + for i in 0..2 { wd[i] = rd[i + shift] as u32; } + } + 1 => { // LOADBU + wd[0] = rd[shift] as u32; + } + 3 => { // STOREW + for i in 0..4 { wd[i] = rd[i] as u32; } + } + 4 => { // STOREH + for i in 0..4 { + if i >= shift && i < 2 + shift { + wd[i] = rd[i - shift] as u32; + } else { + wd[i] = prev[i]; + } + } + } + 5 => { // STOREB + for i in 0..4 { wd[i] = prev[i]; } + wd[shift] = rd[0] as u32; + } + _ => {} // LOADB, LOADH (sign extension) — TODO + } + wd +} + +/// Run ALU operation on 4-byte limbs, returning the result. +fn run_alu_byte(opcode: u8, b: &[u8], c: &[u8]) -> [u8; 4] { + let mut a = [0u8; 4]; + match opcode { + 0 => { + // ADD + let mut carry = 0u32; + for i in 0..4 { + let sum = b[i] as u32 + c[i] as u32 + carry; + a[i] = (sum & 0xFF) as u8; + carry = sum >> 8; + } + } + 1 => { + // SUB + let mut borrow = 0u32; + for i in 0..4 { + let rhs = c[i] as u32 + borrow; + if (b[i] as u32) >= rhs { + a[i] = (b[i] as u32 - rhs) as u8; + borrow = 0; + } else { + a[i] = (256 + b[i] as u32 - rhs) as u8; + borrow = 1; + } + } + } + 2 => { + // XOR + for i in 0..4 { + a[i] = b[i] ^ c[i]; + } + } + 3 => { + // OR + for i in 0..4 { + a[i] = b[i] | c[i]; + } + } + 4 => { + // AND + for i in 0..4 { + a[i] = b[i] & c[i]; + } + } + _ => {} + } + a +} + +/// Arena type determines how records are laid out in memory. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ArenaType { + /// CPU arena: core starts at `adapter_cols_width * sizeof(F)`. + Matrix, + /// GPU arena: core starts at `aligned_adapter_record_size`. + Dense, +} + +/// Core byte offsets per AIR type per arena type. +fn core_byte_offset(air: &str, arena: ArenaType) -> usize { + match (air, arena) { + ("BaseAlu", ArenaType::Matrix) => 19 * 4, // adapter_cols=19 + ("BaseAlu", ArenaType::Dense) => 40, // aligned_adapter_size=40 + ("Shift", ArenaType::Matrix) => 19 * 4, // same adapter as BaseAlu + ("Shift", ArenaType::Dense) => 40, + ("LessThan", ArenaType::Matrix) => 19 * 4, // shares Rv32BaseAluAdapter + ("LessThan", ArenaType::Dense) => 40, + ("LoadStore", ArenaType::Matrix) => 23 * 4, // adapter_cols=23 + ("LoadStore", ArenaType::Dense) => 36, // aligned_adapter_size=36 + ("BranchEqual", ArenaType::Matrix) => 10 * 4, // adapter_cols=10 + ("BranchEqual", ArenaType::Dense) => 24, // aligned_adapter_size=24 + ("BranchLt", ArenaType::Matrix) => 10 * 4, // shares Rv32BranchAdapter + ("BranchLt", ArenaType::Dense) => 24, + ("Auipc", ArenaType::Matrix) => 10 * 4, // Rv32RdWriteAdapter (10 cols) + ("Auipc", ArenaType::Dense) => 20, // 5 u32s in adapter record + ("Jalr", ArenaType::Matrix) => 15 * 4, // Rv32JalrAdapter (15 cols) + ("Jalr", ArenaType::Dense) => 28, // 7 u32s in adapter record + _ => panic!("Unknown AIR type: {air}"), + } +} + +/// Record stride for DenseRecordArena (= aligned_adapter_size + aligned_core_size). +pub fn dense_record_stride(air: &str) -> usize { + match air { + "BaseAlu" | "Shift" | "LessThan" => 52, + "LoadStore" => 60, + "BranchEqual" | "BranchLt" => 40, + "Auipc" => 28, // adapter(20) + core(8: from_pc + imm) + "Jalr" => 44, // adapter(28) + core(16: imm:u16 + 2 pad + from_pc:u32 + rs1_val:u32 + imm_sign:bool + 3 pad) + _ => panic!("Unknown AIR type: {air}"), + } +} + +// ============================================================================ +// Mapping table for BaseAlu (width=36) +// ============================================================================ + +/// Build the mapping table for Rv32BaseAlu (adapter width=19, core width=17, total=36). +/// Uses MatrixRecordArena layout by default (for CPU path). +pub fn base_alu_mapping() -> AirColumnMapping { + base_alu_mapping_for(ArenaType::Matrix) +} + +/// Build BaseAlu mapping for a specific arena type. +pub fn base_alu_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // Adapter record byte offsets (within the adapter section, starting at byte 0) + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rd_ptr: usize = 8; + let rs1_ptr: usize = 12; + let rs2: usize = 16; + let rs2_as: usize = 20; // u8, 1 byte + 3 padding + let reads_aux_0_prev_ts: usize = 24; + let reads_aux_1_prev_ts: usize = 28; + let writes_aux_prev_ts: usize = 32; + let writes_aux_prev_data_0: usize = 36; + + // Core record byte offsets (within the arena row, starting at core offset) + let core = core_byte_offset("BaseAlu", arena); + let b_0: usize = core; + let c_0: usize = core + 4; + let local_opcode: usize = core + 8; + + // Timestamp delta: adapter uses from_timestamp, writes at +2, reads_aux[1] at +1 + let columns = vec![ + // Adapter columns (0-18) + ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }, + ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }, + ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rd_ptr } }, + ColumnMapping { col_index: 3, computation: DirectU32 { record_byte_offset: rs1_ptr } }, + ColumnMapping { col_index: 4, computation: DirectU32 { record_byte_offset: rs2 } }, + ColumnMapping { col_index: 5, computation: DirectU8 { record_byte_offset: rs2_as } }, + // reads_aux[0]: prev_timestamp + decompose(from_ts - prev_ts - 1) + ColumnMapping { col_index: 6, computation: DirectU32 { record_byte_offset: reads_aux_0_prev_ts } }, + ColumnMapping { + col_index: 7, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 0, + prev_ts_byte_offset: reads_aux_0_prev_ts, + limb_index: 0, + }, + }, + ColumnMapping { + col_index: 8, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 0, + prev_ts_byte_offset: reads_aux_0_prev_ts, + limb_index: 1, + }, + }, + // reads_aux[1]: conditional on rs2_as + ColumnMapping { + col_index: 9, + computation: Conditional { + condition_byte_offset: rs2_as, + then_comp: Box::new(DirectU32 { record_byte_offset: reads_aux_1_prev_ts }), + }, + }, + ColumnMapping { + col_index: 10, + computation: Conditional { + condition_byte_offset: rs2_as, + then_comp: Box::new(TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 1, + prev_ts_byte_offset: reads_aux_1_prev_ts, + limb_index: 0, + }), + }, + }, + ColumnMapping { + col_index: 11, + computation: Conditional { + condition_byte_offset: rs2_as, + then_comp: Box::new(TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 1, + prev_ts_byte_offset: reads_aux_1_prev_ts, + limb_index: 1, + }), + }, + }, + // writes_aux: prev_timestamp + decompose(from_ts+2 - prev_ts - 1) + ColumnMapping { col_index: 12, computation: DirectU32 { record_byte_offset: writes_aux_prev_ts } }, + ColumnMapping { + col_index: 13, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 2, + prev_ts_byte_offset: writes_aux_prev_ts, + limb_index: 0, + }, + }, + ColumnMapping { + col_index: 14, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 2, + prev_ts_byte_offset: writes_aux_prev_ts, + limb_index: 1, + }, + }, + // writes_aux.prev_data[0..3] + ColumnMapping { col_index: 15, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 } }, + ColumnMapping { col_index: 16, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 1 } }, + ColumnMapping { col_index: 17, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 2 } }, + ColumnMapping { col_index: 18, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 3 } }, + // Core columns (19-35) + // a[0..3] = run_alu(opcode, b, c) + ColumnMapping { col_index: 19, computation: AluResult { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0, limb_index: 0 } }, + ColumnMapping { col_index: 20, computation: AluResult { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0, limb_index: 1 } }, + ColumnMapping { col_index: 21, computation: AluResult { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0, limb_index: 2 } }, + ColumnMapping { col_index: 22, computation: AluResult { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0, limb_index: 3 } }, + // b[0..3] + ColumnMapping { col_index: 23, computation: DirectU8 { record_byte_offset: b_0 } }, + ColumnMapping { col_index: 24, computation: DirectU8 { record_byte_offset: b_0 + 1 } }, + ColumnMapping { col_index: 25, computation: DirectU8 { record_byte_offset: b_0 + 2 } }, + ColumnMapping { col_index: 26, computation: DirectU8 { record_byte_offset: b_0 + 3 } }, + // c[0..3] + ColumnMapping { col_index: 27, computation: DirectU8 { record_byte_offset: c_0 } }, + ColumnMapping { col_index: 28, computation: DirectU8 { record_byte_offset: c_0 + 1 } }, + ColumnMapping { col_index: 29, computation: DirectU8 { record_byte_offset: c_0 + 2 } }, + ColumnMapping { col_index: 30, computation: DirectU8 { record_byte_offset: c_0 + 3 } }, + // opcode flags + ColumnMapping { col_index: 31, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 0 } }, + ColumnMapping { col_index: 32, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 1 } }, + ColumnMapping { col_index: 33, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 2 } }, + ColumnMapping { col_index: 34, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 3 } }, + ColumnMapping { col_index: 35, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 4 } }, + ]; + + AirColumnMapping { + air_name: "BaseAlu", + width: 36, + record_byte_size: 40 + 12, // adapter(40) + core(12), but need to check padding + columns, + } +} + +// ============================================================================ +// Mapping table for LoadStore (width=41) +// ============================================================================ + +/// Byte offset of the core record start within the arena row for LoadStore. +/// The adapter cols have width 23, so the core bytes start at byte offset 23 * 4 = 92. +const LOADSTORE_CORE_BYTE_OFFSET: usize = 23 * 4; + +/// Build the mapping table for Rv32LoadStore. +pub fn loadstore_mapping() -> AirColumnMapping { + loadstore_mapping_for(ArenaType::Matrix) +} + +pub fn loadstore_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // Adapter record byte offsets + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rs1_ptr: usize = 8; + let rs1_val: usize = 12; // u32 — decomposed into 4 byte limbs for rs1_data + let rs1_aux_prev_ts: usize = 16; + let rd_rs2_ptr: usize = 20; + let read_data_aux_prev_ts: usize = 24; + let imm: usize = 28; // u16 + let imm_sign: usize = 30; // bool/u8 + let mem_as: usize = 31; // u8 + let write_prev_ts: usize = 32; // u32 + + // Core record byte offsets + let core = core_byte_offset("LoadStore", arena); + let local_opcode: usize = core; + let shift_amount: usize = core + 1; + let read_data: usize = core + 2; // u8[4] + let prev_data: usize = core + 8; // u32[4] (after 2 bytes padding at core+6) + + let columns = vec![ + // Adapter columns (0-22) + ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }, + ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }, + ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rs1_ptr } }, + // rs1_data[0..3]: byte decomposition of rs1_val + ColumnMapping { col_index: 3, computation: DirectU8 { record_byte_offset: rs1_val } }, + ColumnMapping { col_index: 4, computation: DirectU8 { record_byte_offset: rs1_val + 1 } }, + ColumnMapping { col_index: 5, computation: DirectU8 { record_byte_offset: rs1_val + 2 } }, + ColumnMapping { col_index: 6, computation: DirectU8 { record_byte_offset: rs1_val + 3 } }, + // rs1_aux_cols + ColumnMapping { col_index: 7, computation: DirectU32 { record_byte_offset: rs1_aux_prev_ts } }, + ColumnMapping { + col_index: 8, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 0, + prev_ts_byte_offset: rs1_aux_prev_ts, + limb_index: 0, + }, + }, + ColumnMapping { + col_index: 9, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 0, + prev_ts_byte_offset: rs1_aux_prev_ts, + limb_index: 1, + }, + }, + // rd_rs2_ptr: conditional (UINT32_MAX means no write => 0) + ColumnMapping { + col_index: 10, + computation: LoadStoreRdRs2Ptr { rd_rs2_ptr_byte_offset: rd_rs2_ptr }, + }, + // read_data_aux + ColumnMapping { col_index: 11, computation: DirectU32 { record_byte_offset: read_data_aux_prev_ts } }, + ColumnMapping { + col_index: 12, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 1, + prev_ts_byte_offset: read_data_aux_prev_ts, + limb_index: 0, + }, + }, + ColumnMapping { + col_index: 13, + computation: TimestampDecomp { + curr_ts_byte_offset: from_timestamp, + curr_ts_delta: 1, + prev_ts_byte_offset: read_data_aux_prev_ts, + limb_index: 1, + }, + }, + // imm (u16) + ColumnMapping { col_index: 14, computation: DirectU16 { record_byte_offset: imm } }, + // imm_sign + ColumnMapping { col_index: 15, computation: DirectU8 { record_byte_offset: imm_sign } }, + // mem_ptr_limbs: computed from rs1_val + sign_extend(imm, imm_sign) + ColumnMapping { + col_index: 16, + computation: PointerLimb { + val_byte_offset: rs1_val, + imm_byte_offset: imm, + imm_sign_byte_offset: imm_sign, + limb_index: 0, + }, + }, + ColumnMapping { + col_index: 17, + computation: PointerLimb { + val_byte_offset: rs1_val, + imm_byte_offset: imm, + imm_sign_byte_offset: imm_sign, + limb_index: 1, + }, + }, + // mem_as + ColumnMapping { col_index: 18, computation: DirectU8 { record_byte_offset: mem_as } }, + // write_base_aux: conditional on needs_write + ColumnMapping { + col_index: 19, + computation: LoadStoreWriteAuxPrevTs { + write_prev_ts_byte_offset: write_prev_ts, + rd_rs2_ptr_byte_offset: rd_rs2_ptr, + }, + }, + ColumnMapping { + col_index: 20, + computation: LoadStoreWriteAuxDecomp { + from_ts_byte_offset: from_timestamp, + write_prev_ts_byte_offset: write_prev_ts, + rd_rs2_ptr_byte_offset: rd_rs2_ptr, + limb_index: 0, + }, + }, + ColumnMapping { + col_index: 21, + computation: LoadStoreWriteAuxDecomp { + from_ts_byte_offset: from_timestamp, + write_prev_ts_byte_offset: write_prev_ts, + rd_rs2_ptr_byte_offset: rd_rs2_ptr, + limb_index: 1, + }, + }, + // needs_write + ColumnMapping { + col_index: 22, + computation: LoadStoreNeedsWrite { rd_rs2_ptr_byte_offset: rd_rs2_ptr }, + }, + // Core columns (23-40) + // flags[0..3]: opcode-dependent + ColumnMapping { col_index: 23, computation: LoadStoreFlag { opcode_byte_offset: local_opcode, shift_byte_offset: shift_amount, flag_index: 0 } }, + ColumnMapping { col_index: 24, computation: LoadStoreFlag { opcode_byte_offset: local_opcode, shift_byte_offset: shift_amount, flag_index: 1 } }, + ColumnMapping { col_index: 25, computation: LoadStoreFlag { opcode_byte_offset: local_opcode, shift_byte_offset: shift_amount, flag_index: 2 } }, + ColumnMapping { col_index: 26, computation: LoadStoreFlag { opcode_byte_offset: local_opcode, shift_byte_offset: shift_amount, flag_index: 3 } }, + // is_valid + ColumnMapping { col_index: 27, computation: Constant(1) }, + // is_load + ColumnMapping { + col_index: 28, + computation: LoadStoreIsLoad { opcode_byte_offset: local_opcode }, + }, + // read_data[0..3] + ColumnMapping { col_index: 29, computation: DirectU8 { record_byte_offset: read_data } }, + ColumnMapping { col_index: 30, computation: DirectU8 { record_byte_offset: read_data + 1 } }, + ColumnMapping { col_index: 31, computation: DirectU8 { record_byte_offset: read_data + 2 } }, + ColumnMapping { col_index: 32, computation: DirectU8 { record_byte_offset: read_data + 3 } }, + // prev_data[0..3]: u32 values from record + ColumnMapping { col_index: 33, computation: DirectU32 { record_byte_offset: prev_data } }, + ColumnMapping { col_index: 34, computation: DirectU32 { record_byte_offset: prev_data + 4 } }, + ColumnMapping { col_index: 35, computation: DirectU32 { record_byte_offset: prev_data + 8 } }, + ColumnMapping { col_index: 36, computation: DirectU32 { record_byte_offset: prev_data + 12 } }, + // write_data[0..3]: computed from opcode, shift, read_data, prev_data + ColumnMapping { + col_index: 37, + computation: LoadStoreWriteData { + opcode_byte_offset: local_opcode, + shift_byte_offset: shift_amount, + read_data_byte_offset: read_data, + prev_data_byte_offset: prev_data, + limb_index: 0, + }, + }, + ColumnMapping { + col_index: 38, + computation: LoadStoreWriteData { + opcode_byte_offset: local_opcode, + shift_byte_offset: shift_amount, + read_data_byte_offset: read_data, + prev_data_byte_offset: prev_data, + limb_index: 1, + }, + }, + ColumnMapping { + col_index: 39, + computation: LoadStoreWriteData { + opcode_byte_offset: local_opcode, + shift_byte_offset: shift_amount, + read_data_byte_offset: read_data, + prev_data_byte_offset: prev_data, + limb_index: 2, + }, + }, + ColumnMapping { + col_index: 40, + computation: LoadStoreWriteData { + opcode_byte_offset: local_opcode, + shift_byte_offset: shift_amount, + read_data_byte_offset: read_data, + prev_data_byte_offset: prev_data, + limb_index: 3, + }, + }, + ]; + + AirColumnMapping { + air_name: "LoadStore", + width: 41, + record_byte_size: 36 + 24, // adapter(36) + core(24) + columns, + } +} + +// ============================================================================ +// Mapping table for Shift (width=53) — uses BaseAlu adapter (width=19) +// ============================================================================ + +/// Build the mapping table for Rv32Shift (adapter=BaseAlu width=19, core=ShiftCore width=34, total=53). +pub fn shift_mapping() -> AirColumnMapping { + shift_mapping_for(ArenaType::Matrix) +} + +pub fn shift_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // Adapter record byte offsets — same as BaseAlu adapter + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rd_ptr: usize = 8; + let rs1_ptr: usize = 12; + let rs2: usize = 16; + let rs2_as: usize = 20; + let reads_aux_0_prev_ts: usize = 24; + let reads_aux_1_prev_ts: usize = 28; + let writes_aux_prev_ts: usize = 32; + let writes_aux_prev_data_0: usize = 36; + + // Core record byte offsets (at core offset = 19*4 = 76) + let core = core_byte_offset("Shift", arena); + let b_0: usize = core; // b[4] at core+0 + let c_0: usize = core + 4; // c[4] at core+4 + let local_opcode: usize = core + 8; // opcode at core+8 + + // ShiftCoreCols layout (starting at col 19): + // a[4], b[4], c[4], sll_flag, srl_flag, sra_flag, + // bit_mul_left, bit_mul_right, b_sign, + // bit_shift_marker[8], limb_shift_marker[4], bit_shift_carry[4] + + let mut columns = Vec::with_capacity(53); + + // Adapter columns (0-18) — identical to BaseAlu + columns.push(ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }); + columns.push(ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }); + columns.push(ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rd_ptr } }); + columns.push(ColumnMapping { col_index: 3, computation: DirectU32 { record_byte_offset: rs1_ptr } }); + columns.push(ColumnMapping { col_index: 4, computation: DirectU32 { record_byte_offset: rs2 } }); + columns.push(ColumnMapping { col_index: 5, computation: DirectU8 { record_byte_offset: rs2_as } }); + columns.push(ColumnMapping { col_index: 6, computation: DirectU32 { record_byte_offset: reads_aux_0_prev_ts } }); + columns.push(ColumnMapping { col_index: 7, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 8, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 1 } }); + columns.push(ColumnMapping { col_index: 9, computation: Conditional { condition_byte_offset: rs2_as, then_comp: Box::new(DirectU32 { record_byte_offset: reads_aux_1_prev_ts }) } }); + columns.push(ColumnMapping { col_index: 10, computation: Conditional { condition_byte_offset: rs2_as, then_comp: Box::new(TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 0 }) } }); + columns.push(ColumnMapping { col_index: 11, computation: Conditional { condition_byte_offset: rs2_as, then_comp: Box::new(TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 1 }) } }); + columns.push(ColumnMapping { col_index: 12, computation: DirectU32 { record_byte_offset: writes_aux_prev_ts } }); + columns.push(ColumnMapping { col_index: 13, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 2, prev_ts_byte_offset: writes_aux_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 14, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 2, prev_ts_byte_offset: writes_aux_prev_ts, limb_index: 1 } }); + columns.push(ColumnMapping { col_index: 15, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 } }); + columns.push(ColumnMapping { col_index: 16, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 1 } }); + columns.push(ColumnMapping { col_index: 17, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 2 } }); + columns.push(ColumnMapping { col_index: 18, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 3 } }); + + // Core columns (19-52): shift-specific + for i in 0..4 { + columns.push(ColumnMapping { col_index: 19 + i, computation: ShiftResult { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0, limb_index: i } }); + } + for i in 0..4 { + columns.push(ColumnMapping { col_index: 23 + i, computation: DirectU8 { record_byte_offset: b_0 + i } }); + } + for i in 0..4 { + columns.push(ColumnMapping { col_index: 27 + i, computation: DirectU8 { record_byte_offset: c_0 + i } }); + } + columns.push(ColumnMapping { col_index: 31, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 0 } }); // sll + columns.push(ColumnMapping { col_index: 32, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 1 } }); // srl + columns.push(ColumnMapping { col_index: 33, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 2 } }); // sra + // bit_multiplier_left, bit_multiplier_right, b_sign + columns.push(ColumnMapping { col_index: 34, computation: ShiftBitMulLeft { opcode_byte_offset: local_opcode, c_byte_offset: c_0 } }); + columns.push(ColumnMapping { col_index: 35, computation: ShiftBitMulRight { opcode_byte_offset: local_opcode, c_byte_offset: c_0 } }); + columns.push(ColumnMapping { col_index: 36, computation: ShiftBSign { opcode_byte_offset: local_opcode, b_byte_offset: b_0 } }); + // bit_shift_marker[8] + for i in 0..8 { + columns.push(ColumnMapping { col_index: 37 + i, computation: ShiftBitMarker { c_byte_offset: c_0, marker_index: i } }); + } + // limb_shift_marker[4] + for i in 0..4 { + columns.push(ColumnMapping { col_index: 45 + i, computation: ShiftLimbMarker { c_byte_offset: c_0, marker_index: i } }); + } + // bit_shift_carry[4] + for i in 0..4 { + columns.push(ColumnMapping { col_index: 49 + i, computation: ShiftBitCarry { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0, limb_index: i } }); + } + + AirColumnMapping { + air_name: "Shift", + width: 53, + record_byte_size: 40 + 12, + columns, + } +} + +// ============================================================================ +// Mapping table for BranchEqual (width=26) +// ============================================================================ + +/// Build the mapping table for BranchEqual (adapter=BranchAdapter width=10, core=BranchEqualCore width=16, total=26). +pub fn branch_equal_mapping() -> AirColumnMapping { + branch_equal_mapping_for(ArenaType::Matrix) +} + +pub fn branch_equal_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // BranchAdapter record byte offsets + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rs1_ptr: usize = 8; + let rs2_ptr: usize = 12; + let reads_aux_0_prev_ts: usize = 16; + let reads_aux_1_prev_ts: usize = 20; + // adapter record = 24 bytes = 6 u32s + + // BranchAdapterCols layout: + // col 0: from_state.pc, col 1: from_state.timestamp + // col 2: rs1_ptr, col 3: rs2_ptr + // col 4: reads_aux_0.prev_timestamp + // col 5-6: reads_aux_0.lt_decomp[0..1] + // col 7: reads_aux_1.prev_timestamp + // col 8-9: reads_aux_1.lt_decomp[0..1] + // adapter width = 10 + + // Core record byte offsets + let core = core_byte_offset("BranchEqual", arena); + let a_0: usize = core; // a[4] at core+0 + let b_0: usize = core + 4; // b[4] at core+4 + let core_imm: usize = core + 8; // imm (u32) at core+8 + let core_opcode: usize = core + 12; // local_opcode (u8) at core+12 + + // BranchEqualCoreCols layout (starting at col 10): + // a[4], b[4], cmp_result, imm, opcode_beq_flag, opcode_bne_flag, diff_inv_marker[4] + // core width = 4+4+1+1+1+1+4 = 16. Total = 10+16 = 26 + + let mut columns = Vec::with_capacity(26); + + // Adapter columns (0-9) + columns.push(ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }); + columns.push(ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }); + columns.push(ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rs1_ptr } }); + columns.push(ColumnMapping { col_index: 3, computation: DirectU32 { record_byte_offset: rs2_ptr } }); + columns.push(ColumnMapping { col_index: 4, computation: DirectU32 { record_byte_offset: reads_aux_0_prev_ts } }); + columns.push(ColumnMapping { col_index: 5, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 6, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 1 } }); + columns.push(ColumnMapping { col_index: 7, computation: DirectU32 { record_byte_offset: reads_aux_1_prev_ts } }); + columns.push(ColumnMapping { col_index: 8, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 9, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 1 } }); + + // Core columns (10-25) + for i in 0..4 { + columns.push(ColumnMapping { col_index: 10 + i, computation: DirectU8 { record_byte_offset: a_0 + i } }); + } + for i in 0..4 { + columns.push(ColumnMapping { col_index: 14 + i, computation: DirectU8 { record_byte_offset: b_0 + i } }); + } + columns.push(ColumnMapping { col_index: 18, computation: BranchEqualCmpResult { a_byte_offset: a_0, b_byte_offset: b_0, opcode_byte_offset: core_opcode } }); + columns.push(ColumnMapping { col_index: 19, computation: DirectU32 { record_byte_offset: core_imm } }); + columns.push(ColumnMapping { col_index: 20, computation: BoolFromOpcode { opcode_byte_offset: core_opcode, expected_opcode: 0 } }); // beq + columns.push(ColumnMapping { col_index: 21, computation: BoolFromOpcode { opcode_byte_offset: core_opcode, expected_opcode: 1 } }); // bne + for i in 0..4 { + columns.push(ColumnMapping { col_index: 22 + i, computation: BranchEqualDiffInvMarker { a_byte_offset: a_0, b_byte_offset: b_0, opcode_byte_offset: core_opcode, marker_index: i } }); + } + + AirColumnMapping { + air_name: "BranchEqual", + width: 26, + record_byte_size: 24 + 16, // adapter(24) + core(16 incl padding) + columns, + } +} + +// ============================================================================ +// Mapping table for LessThan (Rv32BaseAluAdapter, width=37) +// ============================================================================ + +/// Build the mapping table for LessThan (slt/sltu/slti/sltiu). +/// Adapter shared with BaseAlu/Shift (width=19); LessThanCore width=18. +pub fn less_than_mapping() -> AirColumnMapping { + less_than_mapping_for(ArenaType::Matrix) +} + +pub fn less_than_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // Rv32BaseAluAdapter byte offsets (identical layout to BaseAlu/Shift). + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rd_ptr: usize = 8; + let rs1_ptr: usize = 12; + let rs2: usize = 16; + let rs2_as: usize = 20; + let reads_aux_0_prev_ts: usize = 24; + let reads_aux_1_prev_ts: usize = 28; + let writes_aux_prev_ts: usize = 32; + let writes_aux_prev_data_0: usize = 36; + + // LessThanCoreRecord layout: b[4] | c[4] | local_opcode (u8) — same prefix + // as BaseAlu/Shift, but no result limbs (cmp_result is 0/1, derived). + let core = core_byte_offset("LessThan", arena); + let b_0: usize = core; + let c_0: usize = core + 4; + let local_opcode: usize = core + 8; + + // LessThanCoreCols layout (starting at adapter width = 19): + // b[4], c[4], cmp_result, opcode_slt_flag, opcode_sltu_flag, + // b_msb_f, c_msb_f, diff_marker[4], diff_val + // → core width = 4 + 4 + 1 + 2 + 2 + 4 + 1 = 18; total width = 37. + let mut columns = Vec::with_capacity(37); + + // ── Adapter columns (0-18) — identical to BaseAlu adapter ──────────────── + columns.push(ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }); + columns.push(ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }); + columns.push(ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rd_ptr } }); + columns.push(ColumnMapping { col_index: 3, computation: DirectU32 { record_byte_offset: rs1_ptr } }); + columns.push(ColumnMapping { col_index: 4, computation: DirectU32 { record_byte_offset: rs2 } }); + columns.push(ColumnMapping { col_index: 5, computation: DirectU8 { record_byte_offset: rs2_as } }); + columns.push(ColumnMapping { col_index: 6, computation: DirectU32 { record_byte_offset: reads_aux_0_prev_ts } }); + columns.push(ColumnMapping { col_index: 7, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 8, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 1 } }); + columns.push(ColumnMapping { col_index: 9, computation: Conditional { condition_byte_offset: rs2_as, then_comp: Box::new(DirectU32 { record_byte_offset: reads_aux_1_prev_ts }) } }); + columns.push(ColumnMapping { col_index: 10, computation: Conditional { condition_byte_offset: rs2_as, then_comp: Box::new(TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 0 }) } }); + columns.push(ColumnMapping { col_index: 11, computation: Conditional { condition_byte_offset: rs2_as, then_comp: Box::new(TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 1 }) } }); + columns.push(ColumnMapping { col_index: 12, computation: DirectU32 { record_byte_offset: writes_aux_prev_ts } }); + columns.push(ColumnMapping { col_index: 13, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 2, prev_ts_byte_offset: writes_aux_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 14, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 2, prev_ts_byte_offset: writes_aux_prev_ts, limb_index: 1 } }); + columns.push(ColumnMapping { col_index: 15, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 } }); + columns.push(ColumnMapping { col_index: 16, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 1 } }); + columns.push(ColumnMapping { col_index: 17, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 2 } }); + columns.push(ColumnMapping { col_index: 18, computation: DirectU8 { record_byte_offset: writes_aux_prev_data_0 + 3 } }); + + // ── Core columns (19-36) ──────────────────────────────────────────────── + // b[0..3] + for i in 0..4 { + columns.push(ColumnMapping { col_index: 19 + i, computation: DirectU8 { record_byte_offset: b_0 + i } }); + } + // c[0..3] + for i in 0..4 { + columns.push(ColumnMapping { col_index: 23 + i, computation: DirectU8 { record_byte_offset: c_0 + i } }); + } + // cmp_result (col 27): runs run_less_than(opcode, b, c) → cmp. + columns.push(ColumnMapping { col_index: 27, computation: LessThanCmpResult { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0 } }); + // opcode flags: slt=0, sltu=1. + columns.push(ColumnMapping { col_index: 28, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 0 } }); + columns.push(ColumnMapping { col_index: 29, computation: BoolFromOpcode { opcode_byte_offset: local_opcode, expected_opcode: 1 } }); + // b_msb_f, c_msb_f. + columns.push(ColumnMapping { col_index: 30, computation: LessThanBMsbF { opcode_byte_offset: local_opcode, b_byte_offset: b_0 } }); + columns.push(ColumnMapping { col_index: 31, computation: LessThanCMsbF { opcode_byte_offset: local_opcode, c_byte_offset: c_0 } }); + // diff_marker[0..3] — share signature with cmp_result/diff_val for future fusion. + for i in 0..4 { + columns.push(ColumnMapping { col_index: 32 + i, computation: LessThanDiffMarker { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0, marker_index: i } }); + } + // diff_val. + columns.push(ColumnMapping { col_index: 36, computation: LessThanDiffVal { opcode_byte_offset: local_opcode, b_byte_offset: b_0, c_byte_offset: c_0 } }); + + AirColumnMapping { + air_name: "LessThan", + width: 37, + record_byte_size: 40 + 12, // adapter(40) + core(12 with padding) + columns, + } +} + +// ============================================================================ +// Mapping table for BranchLessThan (Rv32BranchAdapter, width=32) +// ============================================================================ + +/// Build the mapping table for BranchLessThan (blt/bltu/bge/bgeu). +/// Adapter shared with BranchEqual (width=10); BranchLtCore width=22. +pub fn branch_lt_mapping() -> AirColumnMapping { + branch_lt_mapping_for(ArenaType::Matrix) +} + +pub fn branch_lt_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // Rv32BranchAdapter byte offsets (same as BranchEqual). + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rs1_ptr: usize = 8; + let rs2_ptr: usize = 12; + let reads_aux_0_prev_ts: usize = 16; + let reads_aux_1_prev_ts: usize = 20; + + // BranchLtCoreRecord layout: a[4] | b[4] | imm (u32) | local_opcode (u8). + let core = core_byte_offset("BranchLt", arena); + let a_0: usize = core; + let b_0: usize = core + 4; + let core_imm: usize = core + 8; + let core_opcode: usize = core + 12; + + // BranchLtCoreCols layout (starting at adapter width = 10): + // a[4], b[4], cmp_result, imm, opcode_blt, opcode_bltu, opcode_bge, + // opcode_bgeu, a_msb_f, b_msb_f, cmp_lt, diff_marker[4], diff_val. + // → core width = 22; total width = 32. + let mut columns = Vec::with_capacity(32); + + // ── Adapter columns (0-9) — identical to BranchEqual adapter ──────────── + columns.push(ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }); + columns.push(ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }); + columns.push(ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rs1_ptr } }); + columns.push(ColumnMapping { col_index: 3, computation: DirectU32 { record_byte_offset: rs2_ptr } }); + columns.push(ColumnMapping { col_index: 4, computation: DirectU32 { record_byte_offset: reads_aux_0_prev_ts } }); + columns.push(ColumnMapping { col_index: 5, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 6, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: reads_aux_0_prev_ts, limb_index: 1 } }); + columns.push(ColumnMapping { col_index: 7, computation: DirectU32 { record_byte_offset: reads_aux_1_prev_ts } }); + columns.push(ColumnMapping { col_index: 8, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 9, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: reads_aux_1_prev_ts, limb_index: 1 } }); + + // ── Core columns (10-31) ──────────────────────────────────────────────── + for i in 0..4 { + columns.push(ColumnMapping { col_index: 10 + i, computation: DirectU8 { record_byte_offset: a_0 + i } }); + } + for i in 0..4 { + columns.push(ColumnMapping { col_index: 14 + i, computation: DirectU8 { record_byte_offset: b_0 + i } }); + } + // cmp_result (branch taken when (cmp_result_internal ^ ge_op) selects taken-side). + columns.push(ColumnMapping { col_index: 18, computation: BranchLtCmpResult { opcode_byte_offset: core_opcode, a_byte_offset: a_0, b_byte_offset: b_0 } }); + // imm (already a u32 in record). + columns.push(ColumnMapping { col_index: 19, computation: DirectU32 { record_byte_offset: core_imm } }); + // opcode flags: BLT=0, BLTU=1, BGE=2, BGEU=3. + columns.push(ColumnMapping { col_index: 20, computation: BoolFromOpcode { opcode_byte_offset: core_opcode, expected_opcode: 0 } }); + columns.push(ColumnMapping { col_index: 21, computation: BoolFromOpcode { opcode_byte_offset: core_opcode, expected_opcode: 1 } }); + columns.push(ColumnMapping { col_index: 22, computation: BoolFromOpcode { opcode_byte_offset: core_opcode, expected_opcode: 2 } }); + columns.push(ColumnMapping { col_index: 23, computation: BoolFromOpcode { opcode_byte_offset: core_opcode, expected_opcode: 3 } }); + // a_msb_f, b_msb_f. + columns.push(ColumnMapping { col_index: 24, computation: BranchLtAMsbF { opcode_byte_offset: core_opcode, a_byte_offset: a_0 } }); + columns.push(ColumnMapping { col_index: 25, computation: BranchLtBMsbF { opcode_byte_offset: core_opcode, b_byte_offset: b_0 } }); + // cmp_lt (independent of opcode direction). + columns.push(ColumnMapping { col_index: 26, computation: BranchLtCmpLt { opcode_byte_offset: core_opcode, a_byte_offset: a_0, b_byte_offset: b_0 } }); + // diff_marker[0..3]. + for i in 0..4 { + columns.push(ColumnMapping { col_index: 27 + i, computation: BranchLtDiffMarker { opcode_byte_offset: core_opcode, a_byte_offset: a_0, b_byte_offset: b_0, marker_index: i } }); + } + // diff_val. + columns.push(ColumnMapping { col_index: 31, computation: BranchLtDiffVal { opcode_byte_offset: core_opcode, a_byte_offset: a_0, b_byte_offset: b_0 } }); + + AirColumnMapping { + air_name: "BranchLt", + width: 32, + record_byte_size: 24 + 16, // adapter(24) + core(16 incl padding) + columns, + } +} + +// ============================================================================ +// Mapping table for AUIPC (Rv32RdWriteAdapter, width=20) +// ============================================================================ + +/// Build the mapping table for AUIPC. +/// Adapter (Rv32RdWriteAdapter) width=10; Rv32AuipcCore width=10. +pub fn auipc_mapping() -> AirColumnMapping { + auipc_mapping_for(ArenaType::Matrix) +} + +pub fn auipc_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // Rv32RdWriteAdapter byte offsets (no reads — only one write). + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rd_ptr: usize = 8; + let rd_aux_prev_ts: usize = 12; + let rd_aux_prev_data_0: usize = 16; // 4 bytes (prev_data[0..4]) + + // Rv32AuipcCoreRecord layout: from_pc (u32) | imm (u32). + let core = core_byte_offset("Auipc", arena); + let core_pc: usize = core; + let core_imm: usize = core + 4; + + // Adapter cols (0-9), then core cols (10-19): + // is_valid, imm_limbs[0..3], pc_limbs[0..2], rd_data[0..4] + let mut columns = Vec::with_capacity(20); + + // ── Adapter (0-9) — Rv32RdWriteAdapter ───────────────────────────────── + columns.push(ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }); + columns.push(ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }); + columns.push(ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rd_ptr } }); + columns.push(ColumnMapping { col_index: 3, computation: DirectU32 { record_byte_offset: rd_aux_prev_ts } }); + columns.push(ColumnMapping { col_index: 4, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: rd_aux_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 5, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: rd_aux_prev_ts, limb_index: 1 } }); + for i in 0..4 { + columns.push(ColumnMapping { col_index: 6 + i, computation: DirectU8 { record_byte_offset: rd_aux_prev_data_0 + i } }); + } + + // ── Core (10-19) ─────────────────────────────────────────────────────── + columns.push(ColumnMapping { col_index: 10, computation: Constant(1) }); // is_valid + // imm_limbs[0..3] — bytes 0..2 of imm (limb 3 is constrained to 0). + for i in 0..3 { + columns.push(ColumnMapping { col_index: 11 + i, computation: DirectU8 { record_byte_offset: core_imm + i } }); + } + // pc_limbs[0..2] — bytes 1..2 of pc (the "middle" two limbs). + for i in 0..2 { + columns.push(ColumnMapping { col_index: 14 + i, computation: DirectU8 { record_byte_offset: core_pc + i + 1 } }); + } + // rd_data[0..4] — (pc + (imm << 8)).to_le_bytes()[i]. + for i in 0..4 { + columns.push(ColumnMapping { col_index: 16 + i, computation: AuipcRdLimb { pc_byte_offset: core_pc, imm_byte_offset: core_imm, limb_index: i } }); + } + + AirColumnMapping { + air_name: "Auipc", + width: 20, + record_byte_size: 20 + 8, // adapter(20) + core(8) + columns, + } +} + +// ============================================================================ +// Mapping table for JALR (Rv32JalrAdapter, width=28) +// ============================================================================ + +/// Build the mapping table for JALR. +/// Adapter (Rv32JalrAdapter) width=15; Rv32JalrCore width=13. +pub fn jalr_mapping() -> AirColumnMapping { + jalr_mapping_for(ArenaType::Matrix) +} + +pub fn jalr_mapping_for(arena: ArenaType) -> AirColumnMapping { + use ColumnComputation::*; + + // Rv32JalrAdapter byte offsets. + let from_pc: usize = 0; + let from_timestamp: usize = 4; + let rs1_ptr: usize = 8; + let rd_ptr: usize = 12; // u32::MAX iff no write + let rs1_aux_prev_ts: usize = 16; + let rd_aux_prev_ts: usize = 20; + let rd_aux_prev_data_0: usize = 24; // 4 bytes + + // Rv32JalrCoreRecord layout: imm (u16) | from_pc (u32) | rs1_val (u32) | imm_sign (bool). + let core = core_byte_offset("Jalr", arena); + let core_imm: usize = core; + let core_from_pc: usize = core + 4; + let core_rs1_val: usize = core + 8; + let core_imm_sign: usize = core + 12; + + // Trace col layout (15 adapter + 13 core = 28): + // adapter: + // 0: from_state.pc, 1: from_state.timestamp, + // 2: rs1_ptr, + // 3: rs1_aux.prev_ts, 4-5: rs1_aux ts_decomp, + // 6: rd_ptr (0 if u32::MAX), + // 7: rd_aux.prev_ts (cond), + // 8-9: rd_aux ts_decomp (cond), + // 10-13: rd_aux.prev_data[0..3] (cond), + // 14: needs_write + // core: + // 15: imm (u16), 16-19: rs1_data[0..3], 20-22: rd_data[0..2], + // 23: is_valid, 24: to_pc_least_sig_bit, 25-26: to_pc_limbs[0..1], + // 27: imm_sign + let mut columns = Vec::with_capacity(28); + + // ── Adapter (0-14) — Rv32JalrAdapter ─────────────────────────────────── + columns.push(ColumnMapping { col_index: 0, computation: DirectU32 { record_byte_offset: from_pc } }); + columns.push(ColumnMapping { col_index: 1, computation: DirectU32 { record_byte_offset: from_timestamp } }); + columns.push(ColumnMapping { col_index: 2, computation: DirectU32 { record_byte_offset: rs1_ptr } }); + // rs1 reads_aux is unconditional. + columns.push(ColumnMapping { col_index: 3, computation: DirectU32 { record_byte_offset: rs1_aux_prev_ts } }); + columns.push(ColumnMapping { col_index: 4, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: rs1_aux_prev_ts, limb_index: 0 } }); + columns.push(ColumnMapping { col_index: 5, computation: TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 0, prev_ts_byte_offset: rs1_aux_prev_ts, limb_index: 1 } }); + // rd_ptr — 0 when u32::MAX (matches LoadStoreRdRs2Ptr semantics). + columns.push(ColumnMapping { col_index: 6, computation: LoadStoreRdRs2Ptr { rd_rs2_ptr_byte_offset: rd_ptr } }); + // rd_aux fields gated on needs_write (rd_ptr != u32::MAX). The chip uses + // `from_timestamp + 1` as the write curr_ts (one read happened first), so + // delta=1. We can't reuse LoadStoreWriteAuxDecomp because that arm + // hardcodes delta=2 (LoadStore's two memory accesses). + columns.push(ColumnMapping { + col_index: 7, + computation: ConditionalNotMaxU32 { + ptr_byte_offset: rd_ptr, + then_comp: Box::new(DirectU32 { record_byte_offset: rd_aux_prev_ts }), + }, + }); + columns.push(ColumnMapping { + col_index: 8, + computation: ConditionalNotMaxU32 { + ptr_byte_offset: rd_ptr, + then_comp: Box::new(TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: rd_aux_prev_ts, limb_index: 0 }), + }, + }); + columns.push(ColumnMapping { + col_index: 9, + computation: ConditionalNotMaxU32 { + ptr_byte_offset: rd_ptr, + then_comp: Box::new(TimestampDecomp { curr_ts_byte_offset: from_timestamp, curr_ts_delta: 1, prev_ts_byte_offset: rd_aux_prev_ts, limb_index: 1 }), + }, + }); + for i in 0..4 { + columns.push(ColumnMapping { + col_index: 10 + i, + computation: ConditionalNotMaxU32 { + ptr_byte_offset: rd_ptr, + then_comp: Box::new(DirectU8 { record_byte_offset: rd_aux_prev_data_0 + i }), + }, + }); + } + // needs_write — 1 if rd_ptr != u32::MAX (matches LoadStoreNeedsWrite). + columns.push(ColumnMapping { col_index: 14, computation: LoadStoreNeedsWrite { rd_rs2_ptr_byte_offset: rd_ptr } }); + + // ── Core (15-27) ─────────────────────────────────────────────────────── + columns.push(ColumnMapping { col_index: 15, computation: DirectU16 { record_byte_offset: core_imm } }); + for i in 0..4 { + columns.push(ColumnMapping { col_index: 16 + i, computation: DirectU8 { record_byte_offset: core_rs1_val + i } }); + } + // rd_data[0..2] = top 3 bytes of (pc + 4) → bytes 1..3. + for i in 0..3 { + columns.push(ColumnMapping { col_index: 20 + i, computation: JalrRdLimb { pc_byte_offset: core_from_pc, limb_index: i } }); + } + columns.push(ColumnMapping { col_index: 23, computation: Constant(1) }); // is_valid + columns.push(ColumnMapping { col_index: 24, computation: JalrToPcLsb { rs1_byte_offset: core_rs1_val, imm_byte_offset: core_imm, imm_sign_byte_offset: core_imm_sign } }); + for i in 0..2 { + columns.push(ColumnMapping { col_index: 25 + i, computation: JalrToPcLimb { rs1_byte_offset: core_rs1_val, imm_byte_offset: core_imm, imm_sign_byte_offset: core_imm_sign, limb_index: i } }); + } + columns.push(ColumnMapping { col_index: 27, computation: DirectU8 { record_byte_offset: core_imm_sign } }); + + AirColumnMapping { + air_name: "Jalr", + width: 28, + record_byte_size: 28 + 16, // adapter(28) + core(16: u16 imm + 2 pad + 2 u32 + bool + 3 pad, aligned to u32) + columns, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_base_alu_mapping_coverage() { + let mapping = base_alu_mapping(); + assert_eq!(mapping.width, 36); + assert_eq!(mapping.columns.len(), 36); + // Verify all column indices are covered + for (i, col) in mapping.columns.iter().enumerate() { + assert_eq!(col.col_index, i, "Column {i} has wrong index"); + } + } + + #[test] + fn test_alu_add() { + let result = run_alu_byte(0, &[10, 20, 30, 40], &[5, 10, 15, 20]); + assert_eq!(result, [15, 30, 45, 60]); + } + + #[test] + fn test_alu_add_carry() { + let result = run_alu_byte(0, &[200, 0, 0, 0], &[100, 0, 0, 0]); + assert_eq!(result, [44, 1, 0, 0]); // 200+100=300, 300&0xFF=44, carry=1 + } + + #[test] + fn test_alu_xor() { + let result = run_alu_byte(2, &[0xFF, 0x0F, 0xAA, 0x55], &[0x0F, 0xF0, 0x55, 0xAA]); + assert_eq!(result, [0xF0, 0xFF, 0xFF, 0xFF]); + } + + #[test] + fn test_alu_sub() { + let result = run_alu_byte(1, &[20, 30, 40, 50], &[5, 10, 15, 20]); + assert_eq!(result, [15, 20, 25, 30]); + } + + use openvm_stark_backend::p3_field::PrimeCharacteristicRing; + + /// Validate the JIT mapping against known trace data from the keccak mock prove debug output. + /// Row 0 of BaseAlu from the actual run: + /// Raw record (u32 LE): [2100252, 2805, 48, 92, 28, 1, 2801, 2785, 2462, 1, 0...0, 2, 0...0] + /// After fill_trace: [2100252, 2805, 48, 92, 28, 1, 2801, 3, 0, 2785, 20, 0, 2462, 344, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0] + #[test] + fn test_base_alu_mapping_against_real_data() { + use openvm_stark_sdk::p3_baby_bear::BabyBear; + + let mapping = base_alu_mapping(); + + // Construct the raw arena row bytes as they appear in memory. + // The arena stores the record struct bytes at the positions corresponding to the row buffer. + // Adapter record: 40 bytes at positions 0..39 + // Core record: 12 bytes at positions 76..87 (core offset = 19 * 4 = 76) + let mut record_bytes = vec![0u8; 36 * 4]; // 36 columns * 4 bytes each + + // Adapter record fields: + record_bytes[0..4].copy_from_slice(&2100252u32.to_le_bytes()); // from_pc + record_bytes[4..8].copy_from_slice(&2805u32.to_le_bytes()); // from_timestamp + record_bytes[8..12].copy_from_slice(&48u32.to_le_bytes()); // rd_ptr + record_bytes[12..16].copy_from_slice(&92u32.to_le_bytes()); // rs1_ptr + record_bytes[16..20].copy_from_slice(&28u32.to_le_bytes()); // rs2 + record_bytes[20] = 1; // rs2_as (u8) + record_bytes[24..28].copy_from_slice(&2801u32.to_le_bytes()); // reads_aux[0].prev_ts + record_bytes[28..32].copy_from_slice(&2785u32.to_le_bytes()); // reads_aux[1].prev_ts + record_bytes[32..36].copy_from_slice(&2462u32.to_le_bytes()); // writes_aux.prev_ts + record_bytes[36] = 1; // writes_aux.prev_data[0] + record_bytes[37] = 0; // writes_aux.prev_data[1] + record_bytes[38] = 0; // writes_aux.prev_data[2] + record_bytes[39] = 0; // writes_aux.prev_data[3] + + // Core record at offset 76: + let core_off = 76; + record_bytes[core_off] = 0; // b[0] + record_bytes[core_off + 1] = 0; // b[1] + record_bytes[core_off + 2] = 0; // b[2] + record_bytes[core_off + 3] = 0; // b[3] + record_bytes[core_off + 4] = 0; // c[0] + record_bytes[core_off + 5] = 0; // c[1] + record_bytes[core_off + 6] = 0; // c[2] + record_bytes[core_off + 7] = 0; // c[3] + record_bytes[core_off + 8] = 2; // local_opcode = XOR + + // Expected output from fill_trace_row (as u32): + let expected: Vec = vec![ + 2100252, 2805, 48, 92, 28, 1, 2801, 3, 0, 2785, 20, 0, 2462, 344, 0, 1, 0, 0, 0, + // core: a=[0,0,0,0], b=[0,0,0,0], c=[0,0,0,0], flags=[0,0,1,0,0] + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, + ]; + + // Evaluate each column using the mapping + // Use range_max_bits from the actual run. Looking at the debug output, + // timestamp decomp: diff = 2805 - 2801 - 1 = 3, lower_decomp[0] = 3, [1] = 0 + // This means range_max_bits is large enough that 3 fits in one limb. + // BabyBear range checker typically uses max_bits = 29 or similar. + // Let's figure it out: decompose(3, max_bits) with limb[0]=3, limb[1]=0 + // means 3 < 2^max_bits, so max_bits >= 2. The actual value doesn't matter + // as long as 3 fits in the lower limb. Let's use 17 (common value). + let range_max_bits = 17u32; + + for col_map in &mapping.columns { + let jit_val: BabyBear = + eval_column(&col_map.computation, &record_bytes, range_max_bits); + let expected_val = BabyBear::from_u32(expected[col_map.col_index]); + assert_eq!( + jit_val, expected_val, + "Mismatch at column {} ({}): JIT={}, expected={}", + col_map.col_index, + match &col_map.computation { + ColumnComputation::DirectU32 { .. } => "DirectU32", + ColumnComputation::DirectU8 { .. } => "DirectU8", + ColumnComputation::TimestampDecomp { .. } => "TimestampDecomp", + ColumnComputation::AluResult { .. } => "AluResult", + ColumnComputation::BoolFromOpcode { .. } => "BoolFromOpcode", + ColumnComputation::Conditional { .. } => "Conditional", + ColumnComputation::Constant(_) => "Constant", + _ => "Unsupported", + }, + jit_val.as_canonical_u32(), + expected_val.as_canonical_u32() + ); + } + } +} diff --git a/openvm/src/powdr_extension/trace_generator/mod.rs b/openvm/src/powdr_extension/trace_generator/mod.rs index 753d511f6b..632d6614e8 100644 --- a/openvm/src/powdr_extension/trace_generator/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/mod.rs @@ -1,6 +1,7 @@ pub mod cpu; #[cfg(feature = "cuda")] pub mod cuda; +pub mod jit_mapping; mod common; diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs index a7f1d60615..be0c382dd3 100644 --- a/openvm/src/powdr_extension/vm.rs +++ b/openvm/src/powdr_extension/vm.rs @@ -23,8 +23,8 @@ use openvm_circuit::{ circuit_derive::Chip, }; use openvm_stark_backend::{ - config::{StarkGenericConfig, Val}, p3_field::{Field, PrimeField32}, + StarkProtocolConfig, Val, }; use serde::{Deserialize, Serialize}; @@ -119,7 +119,7 @@ impl VmExecutionExtension for PowdrExtension VmCircuitExtension for PowdrExtension, ISA> where - SC: StarkGenericConfig, + SC: StarkProtocolConfig, Val: PrimeField32, { fn extend_circuit(&self, inventory: &mut AirInventory) -> Result<(), AirInventoryError> { diff --git a/openvm/src/trace_generation.rs b/openvm/src/trace_generation.rs index 291454190c..8be57c6a3a 100644 --- a/openvm/src/trace_generation.rs +++ b/openvm/src/trace_generation.rs @@ -3,22 +3,16 @@ use crate::SpecializedConfigCpuBuilder; use crate::{isa::OpenVmISA, program::CompiledProgram, SpecializedConfig}; use openvm_circuit::arch::{ execution_mode::Segment, Executor, MeteredExecutor, PreflightExecutionOutput, - PreflightExecutor, VirtualMachine, VmBuilder, VmCircuitConfig, VmExecutionConfig, VmInstance, + PreflightExecutor, VirtualMachine, VmBuilder, VmExecutionConfig, }; -use openvm_native_circuit::NativeConfig; use openvm_sdk::{ - config::{AppConfig, DEFAULT_APP_LOG_BLOWUP}, - prover::vm::new_local_prover, + config::{AggregationSystemParams, AppConfig}, GenericSdk, StdIn, }; -use openvm_stark_backend::config::Val; -use openvm_stark_backend::{keygen::types::MultiStarkProvingKey, prover::types::ProvingContext}; -use openvm_stark_sdk::{ - config::{ - baby_bear_poseidon2::BabyBearPoseidon2Engine as CpuBabyBearPoseidon2Engine, FriParameters, - }, - engine::{StarkEngine, StarkFriEngine}, -}; +use openvm_stark_backend::StarkEngine; +use openvm_stark_backend::Val; +use openvm_stark_backend::{keygen::types::MultiStarkProvingKey, prover::ProvingContext}; +use openvm_stark_sdk::config::{app_params_with_100_bits_security, MAX_APP_LOG_STACKED_HEIGHT}; use tracing::info_span; use crate::BabyBearSC; @@ -34,9 +28,9 @@ use crate::SpecializedConfigCpuBuilder as SpecializedConfigBuilder; use crate::SpecializedConfigGpuBuilder as SpecializedConfigBuilder; #[cfg(feature = "cuda")] -use openvm_cuda_backend::engine::GpuBabyBearPoseidon2Engine as BabyBearPoseidon2Engine; +use openvm_cuda_backend::BabyBearPoseidon2GpuEngine as BabyBearPoseidon2CpuEngine; #[cfg(not(feature = "cuda"))] -use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine; +use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2CpuEngine; /// Given a program and input, generates the trace segment by segment and calls the provided /// callback with the VM, proving key, and proving context (containing the trace) for each segment. @@ -45,13 +39,14 @@ pub fn do_with_trace( inputs: StdIn, callback: impl FnMut( usize, - &VirtualMachine>, + &VirtualMachine>, &MultiStarkProvingKey, - ProvingContext<::PB>, + ProvingContext<::PB>, ), ) -> Result<(), Box> { - let sdk = PowdrSdk::new(create_app_config(program))?; - do_with_trace_with_sdk::, _>( + let (app_config, agg_params) = create_app_config(program); + let sdk = PowdrSdk::new(app_config, agg_params)?; + do_with_trace_with_sdk::>( program, inputs, sdk, callback, ) } @@ -62,21 +57,29 @@ pub fn do_with_cpu_trace( inputs: StdIn, callback: impl FnMut( usize, - &VirtualMachine>, + &VirtualMachine< + openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2CpuEngine, + SpecializedConfigCpuBuilder, + >, &MultiStarkProvingKey, - ProvingContext<::PB>, + ProvingContext< + ::PB, + >, ), ) -> Result<(), Box> { - let sdk = PowdrSdkCpu::new(create_app_config(program))?; - do_with_trace_with_sdk::, _>( - program, inputs, sdk, callback, - ) + let (app_config, agg_params) = create_app_config(program); + let sdk = PowdrSdkCpu::new(app_config, agg_params)?; + do_with_trace_with_sdk::< + ISA, + openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2CpuEngine, + SpecializedConfigCpuBuilder, + >(program, inputs, sdk, callback) } -fn do_with_trace_with_sdk( +fn do_with_trace_with_sdk( program: &CompiledProgram, inputs: StdIn, - sdk: GenericSdk, + sdk: GenericSdk, mut callback: impl FnMut( usize, &VirtualMachine, @@ -85,70 +88,65 @@ fn do_with_trace_with_sdk( ), ) -> Result<(), Box> where - E: StarkFriEngine, + E: StarkEngine, VB: VmBuilder + Clone, >>::Executor: Executor> + MeteredExecutor> + PreflightExecutor, VB::RecordArena>, - NB: VmBuilder + Clone, - >>::Executor: - PreflightExecutor, NB::RecordArena>, { - let exe = sdk.convert_to_exe(program.exe.clone())?; - // Build owned vm instance, so we can mutate it later - let vm_builder = sdk.app_vm_builder().clone(); - let vm_pk = sdk.app_pk().app_vm_pk.clone(); - let mut vm_instance: VmInstance<_, _> = new_local_prover(vm_builder, &vm_pk, exe.clone())?; + let app_pk = sdk.app_pk(); + let pk = &*app_pk.app_vm_pk.vm_pk; - vm_instance.reset_state(inputs.clone()); - let metered_ctx = vm_instance.vm.build_metered_ctx(&exe); - let metered_interpreter = vm_instance.vm.metered_interpreter(vm_instance.exe())?; - let (segments, _) = metered_interpreter.execute_metered(inputs.clone(), metered_ctx)?; - let mut state = vm_instance.state_mut().take(); + let mut app_prover = sdk.app_prover(program.exe.clone())?; + let instance = app_prover.instance_mut(); - // Move `vm` and `interpreter` out of `vm_instance` - // (after this, you can't use `vm_instance` anymore). - let mut vm = vm_instance.vm; - let mut interpreter = vm_instance.interpreter; - - // Get reusable inputs for `debug_proving_ctx`, the mock prover API from OVM. - let air_inv = vm.config().create_airs()?; - let pk = air_inv.keygen::(&vm.engine); + // Metered execution to get segment boundaries + let inputs_streams: openvm_circuit::arch::Streams> = inputs.into(); + instance.reset_state(inputs_streams.clone()); + let exe = instance.exe().clone(); + let metered_ctx = instance.vm.build_metered_ctx(&exe); + let metered_interpreter = instance.vm.metered_interpreter(&exe)?; + let (segments, _) = metered_interpreter.execute_metered(inputs_streams, metered_ctx)?; + // For each segment: preflight → generate proving context → callback + let mut state = instance.state_mut().take(); for (seg_idx, segment) in segments.into_iter().enumerate() { - let _segment_span = info_span!("prove_segment", segment = seg_idx).entered(); - // We need a separate span so the metric label includes "segment" from _segment_span - let _prove_span = info_span!("total_proof").entered(); + let _span = info_span!("trace_segment", segment = seg_idx).entered(); let Segment { num_insns, trace_heights, .. } = segment; - let from_state = Option::take(&mut state).unwrap(); - vm.transport_init_memory_to_device(&from_state.memory); + let from_state = state.take().unwrap(); + instance + .vm + .transport_init_memory_to_device(&from_state.memory); let PreflightExecutionOutput { system_records, record_arenas, to_state, - } = vm.execute_preflight( - &mut interpreter, + } = instance.vm.execute_preflight( + &mut instance.interpreter, from_state, Some(num_insns), &trace_heights, )?; state = Some(to_state); - let ctx = vm.generate_proving_ctx(system_records, record_arenas)?; - - callback(seg_idx, &vm, &pk, ctx); + let ctx = instance + .vm + .generate_proving_ctx(system_records, record_arenas)?; + callback(seg_idx, &instance.vm, pk, ctx); } + *instance.state_mut() = state; + Ok(()) } fn create_app_config( program: &CompiledProgram, -) -> AppConfig> { - let app_fri_params = - FriParameters::standard_with_100_bits_conjectured_security(DEFAULT_APP_LOG_BLOWUP); - AppConfig::new(app_fri_params, program.vm_config.clone()) +) -> (AppConfig>, AggregationSystemParams) { + let system_params = app_params_with_100_bits_security(MAX_APP_LOG_STACKED_HEIGHT); + let app_config = AppConfig::new(program.vm_config.clone(), system_params); + (app_config, AggregationSystemParams::default()) }