diff --git a/README.md b/README.md index 935817f..8a8116f 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,20 @@ sasskit analyze kernel.cubin -k MyKernel -t 64 sasskit disassemble kernel.cubin -k MyKernel ``` +Each reforge call uses a private temporary workspace. The input is read-only; +`--work-dir` selects the existing parent directory for temporary files. The best +result is published outside that workspace and remains available after cleanup: + +```bash +sasskit reforge kernel.cubin -k MyKernel -n 300 --work-dir /tmp -o best.cubin +``` + +Existing output files (including input aliases) are never overwritten. Without +`-o`, each call creates `./reforge-results//best.cubin`. +Python callers can use the keyword-only `output_path` and `work_dir` arguments; +`state.best_cubin_path` identifies the published result. A run with no accepted +mutation publishes the baseline if its benchmark is valid. + ## Architecture ``` diff --git a/docs/reforge-workspace-validation.md b/docs/reforge-workspace-validation.md new file mode 100644 index 0000000..76239c5 --- /dev/null +++ b/docs/reforge-workspace-validation.md @@ -0,0 +1,38 @@ +# Reforge workspace validation (SM120) + +Baseline: `kacper-daftcode/sasskit@913c13451285e07e189b8f89440bd848251a50e1`. +Python 3.12.13; tests ran in a dedicated environment under `/home/gongji/0z5a`. +Cubit was built from `1f7a5aa6cb0096223f4054930b58c9c0208251e5`; +pyelftools 0.32. No system environment packages were changed. + +| Check | Baseline | Patched | Speedup | +|---|---|---|---| +| Barrier-synchronized 2-process shared-path regression | Cross-run token overwritten | Pass | N/A: correctness | +| Barrier-synchronized 8-process shared-path regression | Cross-run token overwritten | Pass | N/A: correctness | +| Workspace/CLI/publication suite | New output API absent | 11 passed | N/A | +| Existing suite | 7 skipped | Same 7 skipped | N/A | +| GPU workers sharing one workspace parent | Not run on unsafe baseline | 1 / 2 / 4 / 8 passed | No kernel speedup claimed | + +The CPU baseline reproducer relocates the **two shared filenames only** into a +task-owned directory; it preserves the old shared-path behavior and forces all +writers to reach a barrier before reading. The fixed test uses the actual +`reforge` function, not a replacement lifecycle implementation. + +Tests cover repeated calls, spaces, explicit output collisions, atomic +no-clobber publication, input symlink/hardlink aliases, benchmark exceptions, +KeyboardInterrupt, publication failure, durable returned paths and legacy CLI +arguments. Existing fixture-dependent tests remain skipped because their +original four-kernel test binary is unavailable; they are not counted as passes. + +GPU checks used an explicit `(uint32_t* input, uint32_t* output, uint32_t n)` PTX +fixture, 65,536 elements, 256 blocks × 256 threads, no dynamic shared memory, +and full modulo-2^32 output comparison against `input * 3 + 7`. Each process +verified its actual device UUID against its assigned UUID. Published outputs +were reopened after cleanup and matched the input hash. These runs used +`max_iters=0` and the integrated A1/A2/A3 tree with a test-only numerical runner. +They establish file ownership and numerical survival, not search quality. + +RTX 5090 (SM120), driver 580.82.07, CUDA 13.0.88. The existing fixed-ABI +`sass_test` harness was not treated as a generic numerical oracle. Validation +subprocesses exited naturally; their timeout/termination paths were disabled in +the test adapter. Production timeout behavior was not changed by this PR. diff --git a/src/sasskit/recolor/cli.py b/src/sasskit/recolor/cli.py index 02ecb79..95c65e1 100644 --- a/src/sasskit/recolor/cli.py +++ b/src/sasskit/recolor/cli.py @@ -297,6 +297,8 @@ def cmd_reforge(args: argparse.Namespace) -> int: bench_smem=args.smem, temperature=args.temperature, seed=args.seed, + output_path=args.output, + work_dir=args.work_dir, ) return 0 @@ -761,6 +763,9 @@ def main() -> int: p_reforge.add_argument('--smem', type=int, default=28672) p_reforge.add_argument('--temperature', type=float, default=0.1) p_reforge.add_argument('--seed', type=int, default=42) + p_reforge.add_argument('-o', '--output', + help='Output cubin (no overwrite; default: ./reforge-results//best.cubin)') + p_reforge.add_argument('--work-dir', help='Parent directory for private temporary workspaces') # --- forge commands --- try: diff --git a/src/sasskit/schedule/reforge.py b/src/sasskit/schedule/reforge.py index 41422bd..89c54c1 100644 --- a/src/sasskit/schedule/reforge.py +++ b/src/sasskit/schedule/reforge.py @@ -16,7 +16,11 @@ from __future__ import annotations +import math +import os import random +import shutil +import tempfile import struct import subprocess import sys @@ -71,6 +75,7 @@ class ReforgeState: accepted: int = 0 rejected: int = 0 history: list[tuple[int, str, float]] = field(default_factory=list) + best_cubin_path: str | None = None def _is_memory_load(inst: Instruction) -> bool: @@ -311,6 +316,19 @@ def gpu_bench(cubin_path: str, kernel_name: str, ] +def _publish_best(source: Path, output: Path) -> str: + """Publish complete bytes atomically without replacing an existing result.""" + with tempfile.NamedTemporaryFile(dir=output.parent, prefix='.reforge-', + delete=False) as staging: + staging_path = Path(staging.name) + try: + shutil.copyfile(source, staging_path) + os.link(staging_path, output) + finally: + staging_path.unlink() + return str(output.absolute()) + + def reforge(cubin_path: str, kernel_name: str, max_iters: int = 200, bench_blocks: int = 1, bench_threads: int = 256, @@ -318,7 +336,50 @@ def reforge(cubin_path: str, kernel_name: str, temperature: float = 0.1, cooling: float = 0.995, seed: int = 42, - verbose: bool = True) -> ReforgeState: + verbose: bool = True, *, + output_path: str | Path | None = None, + work_dir: str | Path | None = None) -> ReforgeState: + """Optimize in a private workspace and return a durable, non-overwritten best. + + ``work_dir`` is the parent of a unique temporary directory. By default, + results are published in ``./reforge-results//best.cubin``. + Explicit output parents must already exist. Input files are never modified. + """ + output = Path(output_path).absolute() if output_path is not None else None + if output is not None: + if output.resolve() == Path(cubin_path).resolve() or ( + output.exists() and output.samefile(cubin_path)): + raise ValueError('Output must not alias the input cubin') + if output.exists(): + raise FileExistsError(output) + + with tempfile.TemporaryDirectory(prefix='reforge-', dir=work_dir) as workspace: + current = Path(workspace) / 'current.cubin' + best = Path(workspace) / 'best.cubin' + state = _reforge(cubin_path, kernel_name, max_iters, bench_blocks, + bench_threads, bench_smem, temperature, cooling, seed, + verbose, tmp_path=str(current), best_path=str(best)) + if not math.isfinite(state.best_time_ms) or state.best_time_ms <= 0: + raise ValueError('No valid benchmark result to publish') + if output is None: + results = Path.cwd() / 'reforge-results' + results.mkdir(exist_ok=True) + output = Path(tempfile.mkdtemp(prefix='run-', dir=results)) / 'best.cubin' + state.best_cubin_path = _publish_best(best, output) + if verbose: + print(f" Saved: {state.best_cubin_path}", file=sys.stderr) + return state + + +def _reforge(cubin_path: str, kernel_name: str, + max_iters: int = 200, + bench_blocks: int = 1, bench_threads: int = 256, + bench_smem: int = 28672, + temperature: float = 0.1, + cooling: float = 0.995, + seed: int = 42, + verbose: bool = True, *, + tmp_path: str, best_path: str) -> ReforgeState: """Run the SASS-to-SASS optimization loop. Args: @@ -348,8 +409,6 @@ def reforge(cubin_path: str, kernel_name: str, ) # Baseline benchmark - tmp_path = '/tmp/reforge_current.cubin' - best_path = '/tmp/reforge_best.cubin' cubin.save(tmp_path) cubin.save(best_path) @@ -444,6 +503,5 @@ def reforge(cubin_path: str, kernel_name: str, print(f" Baseline: {baseline:.4f} ms", file=sys.stderr) print(f" Best: {state.best_time_ms:.4f} ms " f"({speedup:.3f}x)", file=sys.stderr) - print(f" Saved: {best_path}", file=sys.stderr) return state diff --git a/tests/test_reforge_workspace.py b/tests/test_reforge_workspace.py new file mode 100644 index 0000000..7f23769 --- /dev/null +++ b/tests/test_reforge_workspace.py @@ -0,0 +1,180 @@ +"""File ownership tests; no CUDA runtime is required.""" +import multiprocessing +from pathlib import Path +from unittest.mock import patch + +import pytest + +from sasskit.schedule import reforge as rf + + +class TokenCubin: + def __init__(self, path): + self.path = Path(path) + self.token = self.path.read_bytes() + + @classmethod + def from_file(cls, path): + return cls(path) + + def get_kernel(self, name): + return name + + def save(self, path): + Path(path).write_bytes(self.token) + + +def run_token(source, barrier=None, **kwargs): + token = Path(source).read_bytes() + saved = [] + + def save(obj, path): + saved.append(Path(path)) + Path(path).write_bytes(obj.token) + + def bench(path, *args): + if barrier is not None: + barrier.wait() + assert Path(path).read_bytes() == token, 'another run overwrote this cubin' + return 1.0 + + with patch.object(rf, 'Cubin', TokenCubin), \ + patch.object(TokenCubin, 'save', save), \ + patch.object(rf, 'decode_kernel', return_value=[]), \ + patch.object(rf, 'gpu_bench', bench): + state = rf.reforge(str(source), 'test', max_iters=0, verbose=False, **kwargs) + return state, saved + + +def worker(source, cwd, barrier): + import os + os.chdir(cwd) + state, saved = run_token(source, barrier) + assert Path(state.best_cubin_path).read_bytes() == Path(source).read_bytes() + assert not saved[0].parent.exists() + + +@pytest.mark.parametrize('count', [2, 8]) +def test_concurrent_calls_share_cwd_and_temp_root(tmp_path, count): + ctx = multiprocessing.get_context('spawn') + barrier = ctx.Barrier(count) + processes = [] + for i in range(count): + source = tmp_path / f'input-{i}.cubin' + source.write_bytes(str(i).encode()) + process = ctx.Process(target=worker, args=(source, tmp_path, barrier)) + process.start() + processes.append(process) + for process in processes: + process.join() + assert [p.exitcode for p in processes] == [0] * count + assert len(list((tmp_path / 'reforge-results').glob('*/best.cubin'))) == count + + +def test_repeated_calls_and_spaces(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + source = tmp_path / 'input with spaces.cubin' + source.write_bytes(b'original') + root = tmp_path / 'temporary root' + root.mkdir() + first, first_saved = run_token(source, work_dir=root) + second, second_saved = run_token(source, work_dir=root) + assert first_saved[0].parent != second_saved[0].parent + assert all(p.parent == first_saved[0].parent for p in first_saved) + assert list(root.iterdir()) == [] + for state in (first, second): + assert Path(state.best_cubin_path).read_bytes() == b'original' + assert state.cubin_path == str(source) + assert first.best_cubin_path != second.best_cubin_path + assert source.read_bytes() == b'original' + + +@pytest.mark.parametrize('alias', ['same', 'symlink', 'hardlink']) +def test_input_alias_rejected(tmp_path, alias): + source = tmp_path / 'input' + source.write_bytes(b'original') + output = tmp_path / 'output' + if alias == 'same': + output = source + elif alias == 'symlink': + output.symlink_to(source) + else: + output.hardlink_to(source) + with pytest.raises(ValueError, match='alias'): + run_token(source, output_path=output) + assert source.read_bytes() == b'original' + + +@pytest.mark.parametrize('error', [RuntimeError('benchmark failed'), KeyboardInterrupt()]) +def test_failure_cleans_only_owned_workspace(tmp_path, error): + source = tmp_path / 'input' + source.write_bytes(b'original') + other = tmp_path / 'reforge-other' + other.mkdir() + (other / 'current.cubin').write_bytes(b'other') + with patch.object(rf, 'Cubin', TokenCubin), \ + patch.object(rf, 'decode_kernel', return_value=[]), \ + patch.object(rf, 'gpu_bench', side_effect=error): + with pytest.raises(type(error)): + rf.reforge(str(source), 'test', work_dir=tmp_path, verbose=False) + assert sorted(p.name for p in tmp_path.iterdir()) == ['input', 'reforge-other'] + assert (other / 'current.cubin').read_bytes() == b'other' + + +def test_publish_failure_and_collision(tmp_path): + source = tmp_path / 'input' + source.write_bytes(b'original') + output = tmp_path / 'output' + with patch.object(rf.os, 'link', side_effect=OSError('publication failed')): + with pytest.raises(OSError, match='publication failed'): + run_token(source, output_path=output, work_dir=tmp_path) + assert sorted(p.name for p in tmp_path.iterdir()) == ['input'] + output.write_bytes(b'preexisting') + with pytest.raises(FileExistsError): + run_token(source, output_path=output) + assert output.read_bytes() == b'preexisting' + assert source.read_bytes() == b'original' + + +def publish_worker(source, output, barrier, results): + barrier.wait() + try: + rf._publish_best(Path(source), Path(output)) + results.put('published') + except FileExistsError: + results.put('conflict') + + +def test_atomic_publish_collision(tmp_path): + ctx = multiprocessing.get_context('spawn') + barrier, results = ctx.Barrier(2), ctx.Queue() + output = tmp_path / 'shared-output' + processes = [] + for i in range(2): + source = tmp_path / f'input-{i}' + source.write_bytes(str(i).encode() * 65536) + process = ctx.Process(target=publish_worker, args=(source, output, barrier, results)) + process.start() + processes.append(process) + for process in processes: + process.join() + assert [p.exitcode for p in processes] == [0, 0] + assert sorted(results.get() for _ in processes) == ['conflict', 'published'] + assert output.read_bytes() in (b'0' * 65536, b'1' * 65536) + assert not list(tmp_path.glob('.reforge-*')) + + +def test_cli_defaults_and_options(tmp_path, monkeypatch, capsys): + from sasskit.recolor import cli + source = tmp_path / 'input' + source.write_bytes(b'original') + with patch.object(cli.Cubin, 'from_file') as load, patch.object(rf, 'reforge') as run: + load.return_value.kernels = {'test': object()} + monkeypatch.setattr('sys.argv', ['sasskit', 'reforge', str(source)]) + assert cli.main() == 0 + assert run.call_args.kwargs['output_path'] is None + assert run.call_args.kwargs['work_dir'] is None + monkeypatch.setattr('sys.argv', ['sasskit', 'reforge', '--help']) + with pytest.raises(SystemExit, match='0'): + cli.main() + assert 'reforge-results' in ''.join(capsys.readouterr().out.split())