Add workload sweeps, per-module comparisons, and GPU experiment runners

.github/workflows/optimization_wave2.yml

@@ -0,0 +1,68 @@
+name: LLVM 21 optimization wave 2
+
+on:
+  workflow_dispatch:
+    inputs:
+      workload:
+        description: Workload to repeat (all runs the complete matrix)
+        type: choice
+        default: all
+        options: [all, small, representative, solana, bitcoin, ethereum, shallenge, self_test]
+  push:
+    branches: [poc/portable-ptx-export]
+    paths:
+      - '.github/workflows/optimization_wave2.yml'
+      - 'examples/ptx_export/run_wave2.py'
+      - 'examples/ptx_export/representative-kernels/**'
+
+permissions:
+  contents: read
+
+jobs:
+  compare:
+    name: Compare ${{ matrix.workload }}
+    runs-on: ubuntu-24.04
+    timeout-minutes: ${{ matrix.workload == 'self_test' && 90 || 45 }}
+    strategy:
+      fail-fast: false
+      matrix:
+        workload: ${{ fromJSON(inputs.workload != '' && inputs.workload != 'all' && format('["{0}"]', inputs.workload) || '["small","representative","solana","bitcoin","ethereum","shallenge","self_test"]') }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: DeterminateSystems/nix-installer-action@main
+      - uses: DeterminateSystems/magic-nix-cache-action@v14
+        with:
+          use-flakehub: false
+      - uses: actions/cache/restore@v4
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target/
+          key: wave2-cargo-${{ runner.os }}-llvm21-${{ github.sha }}
+          restore-keys: |
+            ptx-cargo-v1-${{ runner.os }}-${{ runner.arch }}-llvm21-
+      - name: Build the checked-out backend
+        run: |
+          nix develop .#v21 --command cargo build -p rustc_codegen_nvvm --features llvm21 --target-dir target/cuda-builder-codegen
+          rm -rf target/nvptx64-nvidia-cuda
+      - name: Checkout pinned workload
+        uses: actions/checkout@v4
+        with:
+          repository: brandonros/vanity-miner-rs
+          ref: 9791234249fc8cb762c296c4fda4503d2686ff77
+          path: workloads/vanity-miner
+      - name: Compile and compare
+        run: |
+          mkdir -p artifacts
+          git rev-parse HEAD > artifacts/source-commit.txt
+          sha256sum target/cuda-builder-codegen/debug/librustc_codegen_nvvm.so > artifacts/backend-sha256.txt
+          nix develop .#v21 --command rustc -Vv > artifacts/rustc-version.txt
+          cp flake.lock Cargo.lock rust-toolchain.toml artifacts/
+          nix develop .#v21 --command python3 examples/ptx_export/run_wave2.py ${{ matrix.workload }} --out artifacts/${{ matrix.workload }} --miner workloads/vanity-miner
+      - uses: actions/upload-artifact@v4
+        if: ${{ !cancelled() }}
+        with:
+          name: wave2-${{ matrix.workload }}
+          path: artifacts/
+          if-no-files-found: error

.github/workflows/ptx_export.yml

@@ -4,6 +4,19 @@ on:
   pull_request:
     branches: [main, experiment/cuda13.3-llvm21, "stack/**"]
   workflow_dispatch:
+    inputs:
+      cleanup_experiment:
+        description: Replay LLVM 21 handoff IR with opt-in cleanup passes
+        type: boolean
+        default: false
+      optimization_sweep:
+        description: Run the extended CFG, DCE, memory and inlining experiments
+        type: boolean
+        default: false
+      mining_workload:
+        description: Compile pinned Solana mining source and replay memory cleanup
+        type: boolean
+        default: false
   push:
     branches: [poc/portable-ptx-export]
 
@@ -17,23 +30,94 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: DeterminateSystems/nix-installer-action@main
+      - uses: DeterminateSystems/magic-nix-cache-action@v14
+        with:
+          use-flakehub: false
+      - name: Restore Cargo cache
+        id: cargo-cache
+        uses: actions/cache/restore@v4
+        with:
+          # Include all nested builds, notably target/cuda-builder-codegen.
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target/
+          key: ptx-cargo-v1-${{ runner.os }}-${{ runner.arch }}-llvm21-${{ hashFiles('flake.nix', 'flake.lock', 'rust-toolchain.toml', '**/Cargo.toml', '**/Cargo.lock', '.cargo/**', '.github/workflows/ptx_export.yml') }}-${{ github.sha }}
+          restore-keys: |
+            ptx-cargo-v1-${{ runner.os }}-${{ runner.arch }}-llvm21-${{ hashFiles('flake.nix', 'flake.lock', 'rust-toolchain.toml', '**/Cargo.toml', '**/Cargo.lock', '.cargo/**', '.github/workflows/ptx_export.yml') }}-
+      - name: Test codegen inventory parser
+        run: python3 -m unittest discover -s examples/ptx_export -p 'test_*.py'
+      - name: Rebuild the checked-out backend and regenerate device evidence
+        run: |
+          # cuda_builder searches for an existing backend before trying Cargo.
+          # A restored .so alone therefore does not establish source freshness.
+          nix develop .#v21 --command cargo build -p rustc_codegen_nvvm --features llvm21 --target-dir target/cuda-builder-codegen
+          # Device Cargo fingerprints do not track changes to the backend .so;
+          # cached builds also omit the IR side outputs required by this job.
+          # Retain the expensive host/backend dependency cache.
+          rm -rf target/nvptx64-nvidia-cuda
+      - name: Check guarded-select source semantics on CPU
+        run: |
+          nix develop .#v21 --command rustc --edition=2024 --test examples/ptx_export/kernels/src/guarded_select.rs -o /tmp/guarded-select-debug
+          /tmp/guarded-select-debug
+          nix develop .#v21 --command rustc --edition=2024 -O --test examples/ptx_export/kernels/src/guarded_select.rs -o /tmp/guarded-select-release
+          /tmp/guarded-select-release
       - name: Compile Rust kernels without a GPU
-        run: nix develop .#v21 --command cargo run -p ptx_export --features llvm21 -- artifacts/ptx
+        id: export
+        run: |
+          mkdir -p artifacts/ptx
+          printf '%s\n' 'nix develop .#v21 --command cargo run -vv -p ptx_export --features llvm21 -- artifacts/ptx none' > artifacts/ptx/build-command.txt
+          nix develop .#v21 --command cargo run -vv -p ptx_export --features llvm21 -- artifacts/ptx none 2>&1 | tee artifacts/ptx/build.log
+      - name: Verify LLVM 21 handoff IR
+        run: nix develop .#v21 --command opt-21 -passes=verify -disable-output artifacts/ptx/final-module.ll
       - name: Record provenance
         run: |
           git rev-parse HEAD > artifacts/ptx/source-commit.txt
           nix develop .#v21 --command rustc -Vv > artifacts/ptx/rustc-version.txt
-          sha256sum artifacts/ptx/rust_kernels.ptx > artifacts/ptx/SHA256SUMS
-      - name: Check Python regressions
-        run: python3 -B -m unittest discover -s examples/ptx_export -p 'test_*.py'
-      - name: Check guarded-select semantics
-        run: |
-          nix develop .#v21 --command rustc --edition=2024 --test examples/ptx_export/kernels/src/guarded_select.rs -o /tmp/guarded-select
-          /tmp/guarded-select
-      - name: Check default DCE retention
+          nix develop .#v21 --command nvcc --version > artifacts/ptx/nvcc-version.txt
+          cp flake.lock rust-toolchain.toml Cargo.lock artifacts/ptx/
+          sha256sum target/cuda-builder-codegen/debug/librustc_codegen_nvvm.so > artifacts/ptx/backend-sha256.txt
+          # Keep per-crate LLVM IR as well as the linked final-module.ll.
+          mkdir -p artifacts/ptx/llvm-ir
+          find target/nvptx64-nvidia-cuda -name '*.ll' -exec cp --parents '{}' artifacts/ptx/llvm-ir/ \;
+      - name: Assemble and inspect NVIDIA machine code without a GPU
+        run: nix develop .#v21 --command python3 examples/ptx_export/inspect_codegen.py artifacts/ptx
+      - name: Verify default DCE retention and explicit disable
         run: nix develop .#v21 --command python3 examples/ptx_export/check_default_dce.py artifacts/ptx
+      - name: Investigate pre-NVVM cleanup (opt-in)
+        if: ${{ github.event_name == 'workflow_dispatch' && (inputs.cleanup_experiment || inputs.optimization_sweep) }}
+        run: nix develop .#v21 --command python3 examples/ptx_export/replay_cleanup.py artifacts/ptx ${{ inputs.optimization_sweep && '--extended' || '' }}
+      - name: Validate integrated cleanup against replay (opt-in)
+        if: ${{ github.event_name == 'workflow_dispatch' && (inputs.cleanup_experiment || inputs.optimization_sweep) }}
+        run: nix develop .#v21 --command python3 examples/ptx_export/check_integrated_cleanup.py artifacts/ptx
+      - name: Validate per-module cleanup (extended sweep)
+        if: ${{ !cancelled() && steps.export.outcome == 'success' && github.event_name == 'workflow_dispatch' && inputs.optimization_sweep }}
+        run: nix develop .#v21 --command python3 examples/ptx_export/check_module_cleanup.py artifacts/ptx
+      - name: Validate size-oriented builds (extended sweep)
+        if: ${{ !cancelled() && steps.export.outcome == 'success' && github.event_name == 'workflow_dispatch' && inputs.optimization_sweep }}
+        run: nix develop .#v21 --command python3 examples/ptx_export/check_workloads.py artifacts/ptx
+      - name: Checkout pinned mining workload
+        if: ${{ !cancelled() && steps.export.outcome == 'success' && github.event_name == 'workflow_dispatch' && inputs.mining_workload }}
+        uses: actions/checkout@v4
+        with:
+          repository: brandonros/vanity-miner-rs
+          ref: 9791234249fc8cb762c296c4fda4503d2686ff77
+          path: workloads/vanity-miner
+      - name: Compare cleanup on Solana mining kernel
+        if: ${{ !cancelled() && steps.export.outcome == 'success' && github.event_name == 'workflow_dispatch' && inputs.mining_workload }}
+        run: nix develop .#v21 --command python3 examples/ptx_export/check_workloads.py artifacts/ptx --miner workloads/vanity-miner
       - uses: actions/upload-artifact@v4
+        if: ${{ !cancelled() }}
         with:
           name: rust-ptx
           path: artifacts/ptx/
           if-no-files-found: error
+      - name: Save Cargo cache, including after compilation failures
+        if: ${{ !cancelled() && steps.cargo-cache.outcome == 'success' && steps.cargo-cache.outputs.cache-hit != 'true' }}
+        uses: actions/cache/save@v4
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target/
+          key: ${{ steps.cargo-cache.outputs.cache-primary-key }}

Cargo.toml

@@ -54,6 +54,7 @@ members = [
   "examples/ptx_export",
   "examples/ptx_export/kernels",
   "examples/ptx_export/retention-kernels",
+  "examples/ptx_export/representative-kernels",
 
   "samples/introduction/async_api",
   "samples/introduction/async_api/kernels",

examples/ptx_export/check_module_cleanup.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""Compare opt-in per-codegen-unit cleanup with independent LLVM replay."""
+import argparse
+import hashlib
+import json
+import re
+from pathlib import Path
+import subprocess
+import sys
+from optimization_pipelines import SCALAR
+
+
+def canonical_phi(line):
+    # PHI incoming pairs are unordered. Keep each value attached to its block;
+    # this is applied only after LLVM has verified and stripped local names.
+    if ' = phi ' not in line:
+        return line
+    groups = []
+    depth = 0
+    quoted = escaped = False
+    start = 0
+    for index, char in enumerate(line):
+        if not quoted:
+            if char == '[':
+                if depth == 0: start = index
+                depth += 1
+            elif char == ']':
+                depth -= 1
+                if depth == 0:
+                    group = line[start:index + 1]
+                    if re.search(r',\s*%[\w.$-]+\s*\]$', group):
+                        groups.append((start, index + 1, group))
+        if char == '"' and not escaped: quoted = not quoted
+        escaped = char == "\\" and not escaped
+    if len(groups) < 2:
+        return line
+    if any(line[a[1]:b[0]].strip() != ',' for a, b in zip(groups, groups[1:])):
+        return line
+    return line[:groups[0][0]] + ', '.join(sorted(g[2] for g in groups)) + line[groups[-1][1]:]
+
+
+def canonical(ir):
+    # Ignore LLVM printer comments (including predecessor order), but preserve
+    # semicolons inside quoted identifiers, inline assembly and string constants.
+    lines = []
+    for line in ir.splitlines():
+        quoted = escaped = False
+        end = len(line)
+        for index, char in enumerate(line):
+            if char == ';' and not quoted:
+                end = index
+                break
+            if char == '"' and not escaped:
+                quoted = not quoted
+            escaped = char == "\\" and not escaped
+        text = line[:end].rstrip()
+        if text.strip(): lines.append(canonical_phi(text))
+    return '\n'.join(lines)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('artifacts', type=Path)
+    args = parser.parse_args()
+    root = args.artifacts.resolve()
+    results = []
+    for mode in ('module-scalar', 'module-inline'):
+        out = root/'module-cleanup'/mode; out.mkdir(parents=True, exist_ok=True)
+        command = ['cargo', 'run', '-vv', '-p', 'ptx_export', '--features', 'llvm21', '--', str(out), mode]
+        (out/'compiler-command.json').write_text(json.dumps(command, indent=2)+'\n')
+        with (out/'build.log').open('w') as log:
+            subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True)
+        before = sorted((out/'per-module').glob('*.before.ll'))
+        if not before: raise RuntimeError('no per-module before/after evidence')
+        checked = []
+        for path in before:
+            after = path.with_name(path.name.replace('.before.ll', '.after.ll'))
+            replay = path.with_name(path.name.replace('.before.ll', '.replay.ll'))
+            normalized = path.with_name(path.name.replace('.before.ll', '.normalized.ll'))
+            subprocess.run(['opt-21', '-passes='+SCALAR+',strip-nondebug,verify', '-verify-each', '-S', str(path), '-o', str(replay)], check=True)
+            subprocess.run(['opt-21', '-passes=strip-nondebug,verify', '-S', str(after), '-o', str(normalized)], check=True)
+            if canonical(replay.read_text()) != canonical(normalized.read_text()):
+                raise RuntimeError(f'per-module replay mismatch: {path.name}')
+            checked.append({'module':path.name.removesuffix('.before.ll'),
+                            'before_sha256':hashlib.sha256(path.read_bytes()).hexdigest(),
+                            'after_sha256':hashlib.sha256(after.read_bytes()).hexdigest()})
+        if not any(x['module'].startswith('core') for x in checked):
+            raise RuntimeError('missing dependency module coverage')
+        subprocess.run(['opt-21', '-passes=verify', '-disable-output', str(out/'final-module.ll')], check=True)
+        subprocess.run([sys.executable, str(Path(__file__).with_name('inspect_codegen.py')), str(out)], check=True)
+        if mode == 'module-inline':
+            subprocess.run([sys.executable, str(Path(__file__).with_name('check_cleanup_ir.py')),
+                            str(out/'final-module.ll'), '--out', str(out/'host-ir-check')], check=True)
+        results.append({'mode':mode, 'checked_modules':checked})
+        (root/'module-cleanup'/'comparison.json').write_text(json.dumps(results, indent=2)+'\n')
+        print(f'{mode}: {len(checked)} modules match standalone replay; final PTX assembled', flush=True)
+
+
+if __name__ == '__main__':
+    main()

examples/ptx_export/check_workloads.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+"""Build size-oriented reproducers or pinned, unmodified Solana kernel sources."""
+import argparse
+from contextlib import nullcontext
+import json
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+from replay_cleanup import normalized_functions
+from workload import MINER_COMMIT, prepared_miner
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('artifacts', type=Path)
+    parser.add_argument('--miner', type=Path)
+    args = parser.parse_args()
+    root = args.artifacts.resolve()
+    scripts = Path(__file__).resolve().parent
+    context = prepared_miner(args.miner, root/'mining-solana') if args.miner else nullcontext(None)
+    with context as miner:
+        compare(args, root, scripts, miner)
+
+
+def compare(args, root, scripts, miner):
+    if args.miner:
+        out = root/'mining-solana'
+        (out/'workload.json').write_text(json.dumps({'repository':'brandonros/vanity-miner-rs', 'commit':MINER_COMMIT,
+            'features':['solana'], 'source_change':'cuda_std dependency path only; algorithm source unchanged'}, indent=2)+'\n')
+        builds = [('none', out, [str(miner/'kernels'), 'solana'])]
+    else:
+        builds = [(mode, root/'size-builds'/mode, []) for mode in ('size-s', 'size-z')]
+    for mode, out, extra in builds:
+        out.mkdir(parents=True, exist_ok=True)
+        command = ['cargo', 'run', '-vv', '-p', 'ptx_export', '--features', 'llvm21', '--', str(out), mode, *extra]
+        (out/'compiler-command.json').write_text(json.dumps(command, indent=2)+'\n')
+        with (out/'build.log').open('w') as log:
+            subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True)
+        subprocess.run(['opt-21', '-passes=verify', '-disable-output', str(out/'final-module.ll')], check=True)
+        subprocess.run([sys.executable, str(scripts/'inspect_codegen.py'), str(out)], check=True)
+        if not args.miner:
+            subprocess.run([sys.executable, str(scripts/'check_cleanup_ir.py'), str(out/'final-module.ll'),
+                            '--out', str(out/'host-ir-check')], check=True)
+    if args.miner:
+        shutil.copy2(miner/'kernels/Cargo.lock', out/'Cargo.lock.resolved')
+        subprocess.run([sys.executable, str(scripts/'replay_cleanup.py'), str(out), '--extended', '--only',
+                        'baseline,dce-only,inline-only,inline-cleanup,memory-early-cse,memory-gvn,memory-stores,memory-combined'], check=True)
+        results = []
+        for mode, replay in [('dce', 'dce-only'), ('inline-scalar', 'inline-only')]:
+            dest = out/'integrated-cleanup'/mode
+            dest.mkdir(parents=True, exist_ok=True)
+            command = ['cargo', 'run', '-vv', '-p', 'ptx_export', '--features', 'llvm21', '--',
+                       str(dest), mode, str(miner/'kernels'), 'solana']
+            (dest/'compiler-command.json').write_text(json.dumps(command, indent=2)+'\n')
+            with (dest/'build.log').open('w') as log:
+                subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True)
+            subprocess.run(['opt-21', '-passes=verify', '-disable-output', str(dest/'final-module.ll')], check=True)
+            actual = normalized_functions((dest/'rust_kernels.ptx').read_text())
+            expected = normalized_functions((out/'cleanup-experiment'/replay/'rust_kernels.ptx').read_text())
+            if 'kernel_find_solana_vanity_private_key' not in actual or actual != expected:
+                raise RuntimeError(f'Solana {mode}: integrated cleanup differs from replay')
+            subprocess.run([sys.executable, str(scripts/'inspect_codegen.py'), str(dest)], check=True)
+            results.append({'mode':mode, 'matches_replay_function_bodies':True})
+            (out/'integrated-cleanup/comparison.json').write_text(json.dumps(results, indent=2)+'\n')
+
+
+if __name__ == '__main__':
+    main()