Add bidirectional scan (fast and correctness)

.github/workflows/bench-baseline.yml

@@ -8,6 +8,29 @@ name: bench-baseline
 
 on:
   workflow_dispatch:
+    inputs:
+      run_full_baseline:
+        description: >-
+          Also run the full baseline suite (BASELINE_TEST_PLAN.md). Uncheck to
+          run ONLY the bidirectional long-L sweep, which is usually what you
+          want — the full suite takes far longer.
+        default: true
+        type: boolean
+      long_l_suite:
+        description: "bidirectional long-L suite (sweep-len covers L=128..8192)"
+        default: "sweep-len"
+        type: choice
+        options: [sweep-len, sweep-dim, basic]
+      long_l_compile_max:
+        description: >-
+          Max L for the torch.compile baseline. MEASURED: compile time is LINEAR
+          in L at ~0.26 s per timestep (59s@128, 130s@512, 245s@1024, 524s@2048;
+          reproduced across two runs). So 2048 costs ~9 min of compiling, 4096
+          ~18 min, and 8192 ~36 min — a full sweep to 8192 is ~75 min. At 8192
+          inductor may also OOM building the unrolled graph, which is itself the
+          finding. Set 0 to skip torch.compile entirely.
+        default: "2048"
+        type: string
   push:
     branches: [main]
     paths:
@@ -19,6 +42,8 @@ permissions:
 jobs:
   baseline:
     name: baseline (linux-arm64)
+    # Unchanged on push; on dispatch, opt out to run only the long-L sweep.
+    if: github.event_name != 'workflow_dispatch' || inputs.run_full_baseline
     runs-on: ubuntu-24.04-arm
     timeout-minutes: 120
     steps:
@@ -65,3 +90,87 @@ jobs:
           cat comment.md >> "$GITHUB_STEP_SUMMARY"
           gh api "repos/${{ github.repository }}/commits/${{ github.sha }}/comments" \
             -f body="$(cat comment.md)" --silent || true
+
+  # Bidirectional scan at long sequence lengths — the regime the applications
+  # actually live in (genomics @131k, long audio, multi-hour ECG), and the one
+  # per-PR CI cannot afford. Dispatch-only: torch.compile unrolls the recurrence
+  # into an L-step graph, so the compile step is the entire cost (~8 min at
+  # L<=2048, ~20 min including 8192) and may OOM outright at 8192.
+  #
+  # BOTH outcomes are results worth having:
+  #   - it compiles  -> the vs-torch.compile ratio at the lengths that matter
+  #     (the trend at 128->512 says it should exceed 4.67x)
+  #   - it does not  -> a hard limit on the baseline itself, at exactly the
+  #     sequence lengths our headline applications need
+  long-l:
+    name: bidirectional long-L (linux-arm64)
+    if: github.event_name == 'workflow_dispatch'
+    runs-on: ubuntu-24.04-arm
+    # Compile cost is ~0.26 s/timestep and LINEAR in L, so a full sweep with
+    # compile_max_len=8192 needs ~70 min of inductor alone (35 min for L=8192
+    # by itself). 150 gives real headroom: if the job is KILLED we want to know
+    # it was an OOM, not that we clipped the timeout — those are different
+    # findings and only one of them is interesting.
+    timeout-minutes: 150
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+        with:
+          workspaces: kernel
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Build FFI cdylib
+        run: cargo build --release -p arm-scan-ffi
+        working-directory: kernel
+      - name: Install python deps
+        run: python3 -m pip install --quiet numpy torch
+
+      # Correctness gates speed (CLAUDE.md): never benchmark a broken kernel.
+      - name: Bidirectional correctness through the kernel
+        run: python3 tests/check_bidirectional.py
+
+      - name: Bidirectional long-L sweep
+        env:
+          SUITE: ${{ inputs.long_l_suite }}
+          CMAX: ${{ inputs.long_l_compile_max }}
+        run: |
+          set -o pipefail
+          mkdir -p bench/results
+          EXTRA=""
+          if [ "$CMAX" = "0" ]; then
+            EXTRA="--no-compile"
+          else
+            EXTRA="--compile-max-len $CMAX"
+          fi
+          # The bench flushes its JSON after every shape, so an OOM kill at the
+          # longest L still leaves every shorter shape's results on disk.
+          python3 bench/bench_bidirectional.py \
+            --suite "$SUITE" $EXTRA --tag ci-arm64-longl \
+            --json bench/results/bidi_longl_ci-arm64.json 2>&1 \
+            | tee bench-bidi-longl.log
+
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: bidirectional-long-l
+          path: |
+            bench/results/bidi_longl_ci-arm64.json
+            bench-bidi-longl.log
+          if-no-files-found: warn
+
+      - name: Job summary
+        if: always()
+        run: |
+          {
+            echo "## Bidirectional scan, long L (linux-arm64, provisional)"
+            echo "suite=\`${{ inputs.long_l_suite }}\` compile_max_len=\`${{ inputs.long_l_compile_max }}\`"
+            echo
+            echo "If a shape is missing, torch.compile was killed building its"
+            echo "unrolled graph — which is itself the finding. Partial results"
+            echo "are preserved in the artifact."
+            echo '```'
+            grep -E "^===|fused ==|^  (ref_|scan_fwd|bidirectional)|^  =>|exp-sharing|^bidirectional \(fused\)|torch.compile COST|^  +[0-9]+ +[0-9.]+s|^  +L|^  L=" bench-bidi-longl.log 2>/dev/null || echo "(sweep did not complete)"
+            echo '```'
+          } >> "$GITHUB_STEP_SUMMARY"

.github/workflows/ci.yml

@@ -50,6 +50,13 @@ jobs:
           python3 -m pip install --quiet numpy \
             || python3 -m pip install --quiet --break-system-packages numpy
           python3 tests/check_ffi.py
+      # Definition check for the bidirectional scan: flip-forward-flip must
+      # equal a backward-in-time recurrence. numpy-only (no torch, no kernel),
+      # so it runs on every platform. Also the executable spec for the future
+      # Rust `reverse` flag — see TOPOLOGY_IMPLEMENTATION_PLAN.md §2.
+      - name: Bidirectional definition check (numpy)
+        working-directory: .
+        run: python3 tests/check_bidirectional_math.py
       - name: Bench ladder (scalar vs NEON)
         if: matrix.name == 'linux-arm64'
         run: |
@@ -89,7 +96,9 @@ jobs:
     name: bench-op (linux-arm64)
     needs: test
     runs-on: ubuntu-24.04-arm
-    timeout-minutes: 30
+    # 45, not 30: the bidirectional sweep now compiles the torch.compile baseline
+    # up to L=2048 (~12 min of inductor), on top of the torch install and build.
+    timeout-minutes: 45
     steps:
       - uses: actions/checkout@v4
       - uses: dtolnay/rust-toolchain@stable
@@ -104,6 +113,57 @@ jobs:
       - name: Install torch (CPU aarch64)
         working-directory: .
         run: python3 -m pip install --quiet numpy torch
+      # Fast consolidated sanity: both topologies run, are correct, and beat
+      # native torch — one clear pass/fail before the fuller benchmarks.
+      - name: Smoke check (unidirectional + bidirectional vs torch)
+        working-directory: .
+        run: python3 bench/smoke_topologies.py
+      # Correctness gates speed (CLAUDE.md): the bidirectional wrapper is
+      # checked against the vendored f64 reference through the real kernel
+      # before anything here is benchmarked.
+      - name: Bidirectional correctness through the kernel
+        working-directory: .
+        run: python3 tests/check_bidirectional.py
+      # Bidirectional scan vs stock PyTorch (eager + torch.compile), plus the
+      # internal fused-vs-flip diagnostic. Provisional (shared runner); headline
+      # numbers need a dedicated Arm host per bench/README.md.
+      # Full sweep-len (L = 128 .. 8192) rather than --quick's 128/512, because
+      # the interesting behaviour is at long L: torch.compile scales linearly in
+      # L while the kernel scales sub-linearly, so the gap WIDENS (3.63x at
+      # L=128 -> 4.67x at L=512). Two short shapes could not show that.
+      #
+      # Compilation is the entire cost (~L^0.6: 59s@128, 134s@512), so it is
+      # capped at 2048 — roughly 12 min. L=4096/8192 still get measured against
+      # eager and against the kernel, just without a torch.compile row. Raising
+      # the cap risks inductor OOMing on the unrolled graph; that experiment
+      # belongs in the dispatchable bench-baseline workflow, not on every push.
+      - name: Bidirectional benchmark (vs eager / torch.compile)
+        working-directory: .
+        timeout-minutes: 25
+        run: |
+          set -o pipefail
+          python3 bench/bench_bidirectional.py --suite sweep-len \
+            --compile-max-len 2048 --reps 5 --warmup 1 \
+            --json bench-bidi.json 2>&1 | tee bench-bidi.log
+      - name: Publish bidirectional result
+        if: always()
+        working-directory: .
+        run: |
+          {
+            echo "## Bidirectional scan (linux-arm64 CI runner)"
+            echo "vs-baseline rows are the result. The internal fusion win is small and is NOT a headline — see BIDIRECTIONAL_LOG.md."
+            echo '```'
+            grep -E "^===|fused ==|^  (ref_|scan_fwd|bidirectional)|^  =>|exp-sharing|^bidirectional \(fused\)|torch.compile COST|^  +[0-9]+ +[0-9.]+s|^  +L|^  L=" bench-bidi.log 2>/dev/null || echo "(bench did not run)"
+            echo '```'
+          } >> "$GITHUB_STEP_SUMMARY"
+      # The bench flushes JSON after every shape, so this survives even if a
+      # long-L shape is killed mid-sweep.
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: bidirectional-bench-arm64
+          path: bench-bidi.json
+          if-no-files-found: warn
       - name: Op-level benchmark (kernel vs eager vs torch.compile)
         working-directory: .
         run: |