fix(jit): harden branch-heavy ARM64 traces by siyul-park · Pull Request #137 · siyul-park/minivm

siyul-park · 2026-07-10T00:47:01Z

What changed

harden ARM64 trace spilling around branches, internal calls, entry labels, loops, and terminal mutations
validate and relax out-of-range ARM64 branches with clean threaded fallback
remove warm-path tracer/profiler allocation and lock churn
make constant deduplication safe for uncomparable typed-array constants
add regression, range-boundary, allocation, and large-branch coverage

Why

Branch-heavy LightGBM-style batch traces could spill an F64 value over the Go return address when control flow skipped spill-frame setup or restoration. Large generated programs also amplified profiler growth and tracer coordination overhead.

Validation

go test -race ./...
go vet ./...
go test -race -run '^TestInterpreter_JITArraySetAfterBranchyCallsInLoop$' -count=100 ./interp
git diff --check

Summary by CodeRabbit

New Features
- ARM64 JIT compilation now supports safer handling of long conditional branches.
- Added an option to disable register spilling during native code generation.
- Improved support for resumable and partial execution traces.
Bug Fixes
- Improved validation of branch targets, including program-end exits.
- Prevented invalid branch offsets from producing corrupt native code.
- Fixed constant deduplication for complex and nil values.
- Improved spill-frame handling across native calls and recursion.
Documentation
- Expanded guidance on JIT behavior, branch ranges, tracing, verification, and allocation benchmarks.

coderabbitai · 2026-07-10T00:47:20Z

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 5ce9bfa7-1ea4-43a7-9275-37820d479546

📥 Commits

Reviewing files that changed from the base of the PR and between 8ca10c3 and 1d6b022.

📒 Files selected for processing (42)

AGENTS.md
analysis/blocks.go
analysis/blocks_test.go
asm/arch.go
asm/arm64/abi.go
asm/arm64/abi_arm64.go
asm/arm64/abi_arm64.s
asm/arm64/abi_arm64_test.go
asm/arm64/abi_stub.go
asm/arm64/arch.go
asm/arm64/encoder.go
asm/arm64/encoder_test.go
asm/arm64/frame.go
asm/arm64/frame_test.go
asm/arm64/relax.go
asm/arm64/relax_test.go
asm/assembler.go
asm/assembler_test.go
asm/rewriter.go
benchmarks/alloc_test.go
docs/architecture.md
docs/benchmarks.md
docs/instruction-set.md
docs/jit-internals.md
docs/pass-system.md
docs/verification.md
interp/interp.go
interp/interp_test.go
interp/jit.go
interp/jit_arm64.go
interp/jit_arm64_test.go
interp/pool_test.go
interp/trace.go
interp/trace_test.go
optimize/optimizer_test.go
prof/collector.go
prof/prof_test.go
program/verify.go
program/verify_test.go
transform/cd.go
transform/cd_test.go
transform/dce.go

📝 Walkthrough

Walkthrough

This change adds ARM64 branch range validation and relaxation, updates native ABI and spill-frame handling, supports top-level virtual exits, changes partial trace and JIT behavior, and adds related tests, benchmarks, profiler growth, documentation, and constant-comparison safety.

Changes

Runtime and control-flow changes

Layer / File(s)	Summary
Virtual exits and branch validation `analysis/blocks.go`, `analysis/blocks_test.go`, `program/verify.go`, `program/verify_test.go`, `transform/dce.go`, `optimize/optimizer_test.go`, `docs/pass-system.md`, `docs/verification.md`	Top-level branches may target the past-the-end virtual exit, while function branches remain invalid; block linking, optimization, and dead-code relocation validate the updated control-flow rules.
ARM64 assembler, ABI, and branch relaxation `asm/arch.go`, `asm/arm64/*`, `asm/assembler.go`, `asm/assembler_test.go`, `AGENTS.md`	ARM64 branch immediates are range-checked, supported conditional branches can be expanded into in-range sequences, native contexts use `unsafe.Pointer`, and spill state is preserved across calls.
Entry labels and spill-frame rewriting `asm/rewriter.go`, `docs/architecture.md`	Entry labels are rebased around generated prologues, and rewriting disables or restores spill handling for incompatible intra-code control flow.
Trace recording and JIT execution `interp/interp.go`, `interp/trace.go`, `interp/jit.go`, `interp/jit_arm64.go`, `interp/*_test.go`, `docs/jit-internals.md`, `docs/instruction-set.md`	Trace capture records partial cuts, observation state is memoized, loop and mutation traces apply spill restrictions, and native compilation falls back to threaded execution for unsupported branches or register pressure.
Profiler growth and allocation benchmarks `prof/*`, `benchmarks/alloc_test.go`, `docs/benchmarks.md`	Profiler slices grow through a shared capacity helper, sparse instruction ranges are tested, and warmup versus steady-state allocation benchmarks are added.
Safe constant deduplication `transform/cd.go`, `transform/cd_test.go`	Constant deduplication avoids comparing interface values with uncomparable dynamic types and covers uncomparable and nil constants.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Interpreter
  participant Tracer
  participant JITLowerer
  participant Assembler
  participant ThreadedExecution
  Interpreter->>Tracer: Capture entry, loop, or bounded trace
  Tracer->>JITLowerer: Return complete or partial trace
  JITLowerer->>Assembler: Build native instructions
  Assembler->>JITLowerer: Return native code or branch/register error
  JITLowerer->>ThreadedExecution: Keep threaded execution on unsupported native code

Possibly related PRs

siyul-park/minivm#126: Updates related optimizer tests for executing and validating optimized control-flow programs.

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch codex/lightgbm-tree-perf

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

codecov · 2026-07-10T00:48:08Z

Codecov Report

❌ Patch coverage is 60.23622% with 101 lines in your changes missing coverage. Please review.
✅ Project coverage is 46.52%. Comparing base (8ca10c3) to head (1d6b022).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
asm/assembler.go	11.11%	40 Missing ⚠️
asm/rewriter.go	27.27%	16 Missing ⚠️
interp/interp.go	52.94%	10 Missing and 6 partials ⚠️
asm/arm64/encoder.go	72.72%	6 Missing ⚠️
asm/arm64/frame.go	33.33%	5 Missing and 1 partial ⚠️
interp/jit.go	0.00%	4 Missing ⚠️
analysis/blocks.go	85.71%	1 Missing and 1 partial ⚠️
asm/arm64/relax.go	96.07%	1 Missing and 1 partial ⚠️
interp/trace.go	88.88%	1 Missing and 1 partial ⚠️
program/verify.go	60.00%	1 Missing and 1 partial ⚠️
... and 4 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #137      +/-   ##
==========================================
+ Coverage   45.59%   46.52%   +0.93%     
==========================================
  Files          77       78       +1     
  Lines       18564    18760     +196     
==========================================
+ Hits         8465     8729     +264     
+ Misses       9125     9029      -96     
- Partials      974     1002      +28

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

* fix(jit): harden branch-heavy ARM64 traces * fix: correctness

fix(jit): harden branch-heavy ARM64 traces

aee1faa

siyul-park mentioned this pull request Jul 10, 2026

perf(predict): batch tree model evaluation siyul-park/tl2g#1

Draft

fix: correctness

1d6b022

siyul-park marked this pull request as ready for review July 10, 2026 15:11

siyul-park merged commit fd5b2e9 into main Jul 10, 2026
6 checks passed

siyul-park added a commit that referenced this pull request Jul 10, 2026

fix(jit): harden branch-heavy ARM64 traces (#137)

f8f174e

* fix(jit): harden branch-heavy ARM64 traces * fix: correctness

This was referenced Jul 12, 2026

refactor(interp): unify plan-based JIT compilation #139

Merged

perf(interp): fold loop branch legs into native back-edges #160

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix(jit): harden branch-heavy ARM64 traces#137

fix(jit): harden branch-heavy ARM64 traces#137
siyul-park merged 2 commits into
mainfrom
codex/lightgbm-tree-perf

siyul-park commented Jul 10, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 10, 2026 •

edited

Loading

Review failed

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

codecov Bot commented Jul 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

siyul-park commented Jul 10, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changed

Why

Validation

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review failed

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

codecov Bot commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

siyul-park commented Jul 10, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 10, 2026 •

edited

Loading

codecov Bot commented Jul 10, 2026 •

edited

Loading