test: add runtime correctness and benchmark coverage

Co-authored-by: wawahejun <hejunlbbc@gmail.com>

tests/__init__.py

@@ -0,0 +1 @@
+

tests/bench_all.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+PYTHON_DIR = ROOT / "python"
+if str(PYTHON_DIR) not in sys.path:
+    sys.path.insert(0, str(PYTHON_DIR))
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+import torch
+
+from operator_runtime import copy, reduce_sum, softmax, vector_add
+from operator_runtime.benchmark import cuda_time_ms
+from operator_runtime.perf_model import (
+    estimate_copy,
+    estimate_reduce_sum,
+    estimate_softmax,
+    estimate_vector_add,
+)
+from operator_runtime.profiler import PerformanceResult
+
+
+def bench_copy(backend: str) -> PerformanceResult:
+    src = torch.randn((1 << 20,), dtype=torch.float16, device="cuda")
+    out = torch.empty_like(src)
+    runtime = cuda_time_ms(lambda: copy(src, backend=backend))
+    torch_ms = cuda_time_ms(lambda: out.copy_(src))
+    bytes_, flops = estimate_copy(src)
+    return PerformanceResult("copy", backend, str(tuple(src.shape)), str(src.dtype), bytes_, flops, runtime, torch_ms)
+
+
+def bench_vector_add(backend: str) -> PerformanceResult:
+    a = torch.randn((1 << 20,), dtype=torch.float16, device="cuda")
+    b = torch.randn_like(a)
+    runtime = cuda_time_ms(lambda: vector_add(a, b, backend=backend))
+    torch_ms = cuda_time_ms(lambda: torch.add(a, b))
+    bytes_, flops = estimate_vector_add(a)
+    return PerformanceResult("vector_add", backend, str(tuple(a.shape)), str(a.dtype), bytes_, flops, runtime, torch_ms)
+
+
+def bench_reduce_sum(backend: str) -> PerformanceResult:
+    src = torch.randn((1024, 1024), dtype=torch.float32, device="cuda")
+    runtime = cuda_time_ms(lambda: reduce_sum(src, dim=1, backend=backend))
+    torch_ms = cuda_time_ms(lambda: torch.sum(src, dim=1))
+    bytes_, flops = estimate_reduce_sum(src)
+    return PerformanceResult("reduce_sum", backend, str(tuple(src.shape)), str(src.dtype), bytes_, flops, runtime, torch_ms)
+
+
+def bench_softmax(backend: str) -> PerformanceResult:
+    src = torch.randn((1024, 1024), dtype=torch.float32, device="cuda")
+    runtime = cuda_time_ms(lambda: softmax(src, dim=1, backend=backend))
+    torch_ms = cuda_time_ms(lambda: torch.softmax(src, dim=1))
+    bytes_, flops = estimate_softmax(src)
+    return PerformanceResult("softmax", backend, str(tuple(src.shape)), str(src.dtype), bytes_, flops, runtime, torch_ms)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", default="nvidia")
+    parser.add_argument("--profile", default=None)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        print("CUDA is required for benchmark", file=sys.stderr)
+        return 2
+
+    rows = [
+        bench_copy(args.backend),
+        bench_vector_add(args.backend),
+        bench_reduce_sum(args.backend),
+        bench_softmax(args.backend),
+    ]
+
+    print("operator backend shape dtype runtime_ms torch_ms speedup GB/s GFLOP/s")
+    for row in rows:
+        speedup = 0.0 if row.speedup is None else row.speedup
+        print(
+            f"{row.operator} {row.backend} {row.shape} {row.dtype} "
+            f"{row.runtime_ms:.4f} {row.torch_ms or 0:.4f} {speedup:.2f} "
+            f"{row.gbytes_per_sec:.2f} {row.gflops_per_sec:.2f}"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+

tests/cases/__init__.py

@@ -0,0 +1 @@
+

tests/cases/copy.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import torch
+
+
+def correctness_cases():
+    return [
+        {"name": "contiguous_1k_fp32", "shape": (1024,), "dtype": torch.float32, "atol": 0, "rtol": 0},
+        {"name": "contiguous_1k_fp16", "shape": (1024,), "dtype": torch.float16, "atol": 0, "rtol": 0},
+    ]
+
+
+def api_error_cases():
+    return [
+        {"name": "shape_mismatch", "shape": (16,), "out_shape": (8,), "dtype": torch.float32},
+        {"name": "non_contiguous", "shape": (4, 4), "dtype": torch.float32},
+    ]
+
+
+def benchmark_cases():
+    return [
+        {"name": "contiguous_1m", "shape": (1 << 20,), "dtype": torch.float16},
+    ]
+

tests/cases/reduce_sum.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import torch
+
+
+def correctness_cases():
+    return [
+        {"name": "rowwise_32x128", "shape": (32, 128), "dtype": torch.float32, "atol": 1e-5, "rtol": 1e-5},
+    ]
+
+
+def api_error_cases():
+    return [
+        {"name": "wrong_dim", "shape": (16, 16), "dtype": torch.float32, "dim": 0},
+        {"name": "wrong_dtype", "shape": (16, 16), "dtype": torch.float16, "dim": 1},
+    ]
+
+
+def benchmark_cases():
+    return [
+        {"name": "rowwise_1024x1024", "shape": (1024, 1024), "dtype": torch.float32},
+    ]
+

tests/cases/softmax.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import torch
+
+
+def correctness_cases():
+    return [
+        {"name": "rowwise_32x128", "shape": (32, 128), "dtype": torch.float32, "atol": 1e-5, "rtol": 1e-5},
+    ]
+
+
+def api_error_cases():
+    return [
+        {"name": "wrong_dim", "shape": (16, 16), "dtype": torch.float32, "dim": 0},
+        {"name": "wrong_dtype", "shape": (16, 16), "dtype": torch.float16, "dim": 1},
+    ]
+
+
+def benchmark_cases():
+    return [
+        {"name": "rowwise_1024x1024", "shape": (1024, 1024), "dtype": torch.float32},
+    ]
+

tests/cases/vector_add.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import torch
+
+
+def correctness_cases():
+    return [
+        {"name": "contiguous_1k_fp32", "shape": (1024,), "dtype": torch.float32, "atol": 1e-6, "rtol": 1e-6},
+        {"name": "contiguous_1k_fp16", "shape": (1024,), "dtype": torch.float16, "atol": 1e-3, "rtol": 1e-3},
+    ]
+
+
+def api_error_cases():
+    return [
+        {"name": "shape_mismatch", "shape": (16,), "other_shape": (8,), "dtype": torch.float32},
+        {"name": "non_contiguous", "shape": (4, 4), "dtype": torch.float32},
+    ]
+
+
+def benchmark_cases():
+    return [
+        {"name": "contiguous_1m", "shape": (1 << 20,), "dtype": torch.float16},
+    ]
+

tests/conftest.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import pytest
+
+ROOT = Path(__file__).resolve().parents[1]
+PYTHON_DIR = ROOT / "python"
+if str(PYTHON_DIR) not in sys.path:
+    sys.path.insert(0, str(PYTHON_DIR))
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+
+def pytest_addoption(parser):
+    parser.addoption("--backend", action="store", default=os.environ.get("CAMP_TEST_BACKEND", "nvidia"))
+
+
+@pytest.fixture
+def backend(request) -> str:
+    return request.config.getoption("--backend")
+

tests/perf_profiles/local_gpu.yaml

@@ -0,0 +1,4 @@
+name: local-gpu
+peak_bandwidth_gb_s: 1000
+notes: "Placeholder profile. Replace with a measured HBM microbenchmark value."
+

tests/test_api_contract.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from operator_runtime import copy_, reduce_sum, softmax, vector_add
+from operator_runtime.testing import require_cuda
+
+
+def test_copy_rejects_shape_mismatch(backend):
+    require_cuda()
+    src = torch.randn((16,), device="cuda")
+    out = torch.empty((8,), device="cuda")
+    with pytest.raises(ValueError):
+        copy_(out, src, backend=backend)
+
+
+def test_vector_add_rejects_shape_mismatch(backend):
+    require_cuda()
+    a = torch.randn((16,), device="cuda")
+    b = torch.randn((8,), device="cuda")
+    with pytest.raises(ValueError):
+        vector_add(a, b, backend=backend)
+
+
+def test_reduce_sum_rejects_wrong_dim(backend):
+    require_cuda()
+    src = torch.randn((16, 16), device="cuda")
+    with pytest.raises(ValueError):
+        reduce_sum(src, dim=0, backend=backend)
+
+
+def test_softmax_rejects_wrong_dtype(backend):
+    require_cuda()
+    src = torch.randn((16, 16), dtype=torch.float16, device="cuda")
+    with pytest.raises(TypeError):
+        softmax(src, dim=1, backend=backend)
+

tests/test_correctness.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from operator_runtime import copy, reduce_sum, softmax, vector_add
+from operator_runtime.testing import assert_close, require_cuda
+from tests.cases import copy as copy_cases
+from tests.cases import reduce_sum as reduce_sum_cases
+from tests.cases import softmax as softmax_cases
+from tests.cases import vector_add as vector_add_cases
+
+
+@pytest.mark.parametrize("case", copy_cases.correctness_cases(), ids=lambda c: c["name"])
+def test_copy_correctness(case, backend):
+    require_cuda()
+    src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
+    out = copy(src, backend=backend)
+    assert_close(out, src, atol=case["atol"], rtol=case["rtol"])
+
+
+@pytest.mark.parametrize("case", vector_add_cases.correctness_cases(), ids=lambda c: c["name"])
+def test_vector_add_correctness(case, backend):
+    require_cuda()
+    a = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
+    b = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
+    out = vector_add(a, b, backend=backend)
+    assert_close(out, a + b, atol=case["atol"], rtol=case["rtol"])
+
+
+@pytest.mark.parametrize("case", reduce_sum_cases.correctness_cases(), ids=lambda c: c["name"])
+def test_reduce_sum_correctness(case, backend):
+    require_cuda()
+    src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
+    out = reduce_sum(src, dim=1, backend=backend)
+    assert_close(out, torch.sum(src, dim=1), atol=case["atol"], rtol=case["rtol"])
+
+
+@pytest.mark.parametrize("case", softmax_cases.correctness_cases(), ids=lambda c: c["name"])
+def test_softmax_correctness(case, backend):
+    require_cuda()
+    src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
+    out = softmax(src, dim=1, backend=backend)
+    assert_close(out, torch.softmax(src, dim=1), atol=case["atol"], rtol=case["rtol"])
+

tests/test_descriptor_lifecycle.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from operator_runtime.ops.vector_add import prepare_vector_add
+from operator_runtime.testing import assert_close, require_cuda
+
+
+def test_prepared_vector_add_reuses_descriptor(backend):
+    require_cuda()
+    if backend != "nvidia":
+        pytest.skip("descriptor lifecycle test targets C ABI backend")
+    a = torch.randn((1024,), device="cuda")
+    b = torch.randn((1024,), device="cuda")
+    out = torch.empty_like(a)
+    prepared = prepare_vector_add(out, a, b, backend=backend)
+    try:
+        prepared.run()
+        prepared.run()
+        assert_close(out, a + b, atol=1e-6, rtol=1e-6)
+    finally:
+        prepared.destroy()
+

tests/test_manifest_contract.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+
+def test_manifest_validator_passes():
+    root = Path(__file__).resolve().parents[1]
+    result = subprocess.run(
+        [
+            sys.executable,
+            "tools/validate_operator_manifest.py",
+            "--ops-root",
+            "ops",
+            "--tests-root",
+            "tests",
+        ],
+        cwd=root,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    assert result.returncode == 0, result.stderr
+