377 lines
11 KiB
Python
377 lines
11 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# Copyright (C) 2026 David (Corey) Herrera
|
|
"""
|
|
Comparison utilities for PyTorch vs Foundry testing.
|
|
|
|
This module provides the infrastructure for comparing PyTorch and Foundry outputs.
|
|
"""
|
|
|
|
import numpy as np
|
|
import subprocess
|
|
import struct
|
|
import os
|
|
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from dataclasses import dataclass, asdict
|
|
from typing import Optional, List, Tuple, Callable
|
|
import sys
|
|
|
|
# Paths
|
|
SUITE_DIR = Path(__file__).parent.parent
|
|
DATA_DIR = SUITE_DIR / "data"
|
|
RESULTS_DIR = SUITE_DIR / "results"
|
|
COMPARE_HELPER = SUITE_DIR.parent.parent / "bin" / "compare_helper"
|
|
|
|
# Ensure directories exist
|
|
DATA_DIR.mkdir(exist_ok=True)
|
|
RESULTS_DIR.mkdir(exist_ok=True)
|
|
|
|
# Data type mappings
|
|
DTYPE_MAP = {
|
|
'float32': 0,
|
|
'float64': 1,
|
|
'int32': 2,
|
|
'int64': 3,
|
|
}
|
|
|
|
NUMPY_DTYPE_MAP = {
|
|
0: np.float32,
|
|
1: np.float64,
|
|
2: np.int32,
|
|
3: np.int64,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class CompareResult:
|
|
"""Result of a comparison test."""
|
|
name: str
|
|
passed: bool
|
|
max_diff: float
|
|
mean_diff: float
|
|
tolerance: float
|
|
input_shape: List[int]
|
|
time_pytorch_ms: float
|
|
time_foundry_ms: float
|
|
error_message: Optional[str] = None
|
|
|
|
|
|
def write_tensor(filepath: Path, data: np.ndarray):
|
|
"""
|
|
Write a tensor to binary file in Foundry format.
|
|
|
|
Format:
|
|
- ndim: int32 (4 bytes)
|
|
- dtype: int32 (4 bytes)
|
|
- reserved: int64 (8 bytes)
|
|
- shape: ndim * int64
|
|
- data: numel * dtype
|
|
"""
|
|
ndim = len(data.shape)
|
|
dtype_code = DTYPE_MAP.get(str(data.dtype), 0)
|
|
|
|
with open(filepath, 'wb') as f:
|
|
# Header
|
|
f.write(struct.pack('<i', ndim))
|
|
f.write(struct.pack('<i', dtype_code))
|
|
f.write(struct.pack('<q', 0)) # reserved
|
|
|
|
# Shape
|
|
for dim in data.shape:
|
|
f.write(struct.pack('<q', dim))
|
|
|
|
# Data (ensure contiguous)
|
|
f.write(np.ascontiguousarray(data).tobytes())
|
|
|
|
|
|
def read_tensor(filepath: Path) -> np.ndarray:
|
|
"""
|
|
Read a tensor from binary file in Foundry format.
|
|
"""
|
|
with open(filepath, 'rb') as f:
|
|
# Header
|
|
ndim = struct.unpack('<i', f.read(4))[0]
|
|
dtype_code = struct.unpack('<i', f.read(4))[0]
|
|
_ = struct.unpack('<q', f.read(8))[0] # reserved
|
|
|
|
# Shape
|
|
shape = []
|
|
for _ in range(ndim):
|
|
shape.append(struct.unpack('<q', f.read(8))[0])
|
|
|
|
# Data
|
|
dtype = NUMPY_DTYPE_MAP.get(dtype_code, np.float32)
|
|
numel = int(np.prod(shape)) if shape else 1
|
|
data = np.frombuffer(f.read(numel * np.dtype(dtype).itemsize), dtype=dtype)
|
|
|
|
return data.reshape(shape) if shape else data
|
|
|
|
|
|
def run_foundry(operation: str, inputs: List[str], output: str,
|
|
params: Optional[dict] = None) -> Tuple[bool, float]:
|
|
"""
|
|
Run a Foundry operation via the comparison helper.
|
|
|
|
Args:
|
|
operation: Name of the operation (e.g., "mish", "matmul")
|
|
inputs: List of input file paths
|
|
output: Output file path
|
|
params: Optional parameters (e.g., {"alpha": 1.0})
|
|
|
|
Returns:
|
|
Tuple of (success, time_ms)
|
|
"""
|
|
if not COMPARE_HELPER.exists():
|
|
raise FileNotFoundError(
|
|
f"Compare helper not found at {COMPARE_HELPER}. "
|
|
"Run 'make compare-helper' first."
|
|
)
|
|
|
|
cmd = [str(COMPARE_HELPER), operation, output] + inputs
|
|
|
|
if params:
|
|
for key, value in params.items():
|
|
cmd.extend([f"--{key}", str(value)])
|
|
|
|
import time
|
|
start = time.perf_counter()
|
|
|
|
try:
|
|
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
|
elapsed = (time.perf_counter() - start) * 1000
|
|
|
|
if result.returncode != 0:
|
|
print(f"Foundry error: {result.stderr}", file=sys.stderr)
|
|
return False, elapsed
|
|
|
|
return True, elapsed
|
|
|
|
except subprocess.TimeoutExpired:
|
|
return False, 30000.0
|
|
except Exception as e:
|
|
print(f"Error running Foundry: {e}", file=sys.stderr)
|
|
return False, 0.0
|
|
|
|
|
|
def compare_tensors(expected: np.ndarray, actual: np.ndarray,
|
|
tolerance: float) -> Tuple[bool, float, float]:
|
|
"""
|
|
Compare two tensors for numerical equivalence.
|
|
|
|
Returns:
|
|
Tuple of (passed, max_diff, mean_diff)
|
|
"""
|
|
# Squeeze both to handle scalar shape differences: () vs (1,)
|
|
expected = np.squeeze(expected)
|
|
actual = np.squeeze(actual)
|
|
|
|
if expected.shape != actual.shape:
|
|
return False, float('inf'), float('inf')
|
|
|
|
diff = np.abs(expected.astype(np.float64) - actual.astype(np.float64))
|
|
max_diff = float(np.max(diff))
|
|
mean_diff = float(np.mean(diff))
|
|
|
|
passed = max_diff <= tolerance
|
|
|
|
return passed, max_diff, mean_diff
|
|
|
|
|
|
class ComparisonTest:
|
|
"""Base class for comparison tests."""
|
|
|
|
def __init__(self, name: str, tolerance: float = 1e-5):
|
|
self.name = name
|
|
self.tolerance = tolerance
|
|
self.results: List[CompareResult] = []
|
|
|
|
def generate_input(self, seed: int = 42) -> Tuple[np.ndarray, ...]:
|
|
"""Generate input tensors. Override in subclass."""
|
|
raise NotImplementedError
|
|
|
|
def pytorch_forward(self, *inputs) -> np.ndarray:
|
|
"""Run PyTorch forward. Override in subclass."""
|
|
raise NotImplementedError
|
|
|
|
def foundry_operation(self) -> str:
|
|
"""Return the Foundry operation name."""
|
|
return self.name
|
|
|
|
def foundry_params(self) -> Optional[dict]:
|
|
"""Return any extra parameters for Foundry."""
|
|
return None
|
|
|
|
def run(self, seed: int = 42, debug: bool = False) -> CompareResult:
|
|
"""Run the comparison test."""
|
|
import time
|
|
|
|
try:
|
|
# Generate inputs
|
|
inputs = self.generate_input(seed)
|
|
|
|
# Run PyTorch
|
|
import torch
|
|
torch_inputs = [torch.from_numpy(x) for x in inputs]
|
|
|
|
start = time.perf_counter()
|
|
expected = self.pytorch_forward(*torch_inputs)
|
|
if isinstance(expected, torch.Tensor):
|
|
expected = expected.detach().numpy()
|
|
pytorch_time = (time.perf_counter() - start) * 1000
|
|
|
|
# Write inputs to files
|
|
input_files = []
|
|
for i, inp in enumerate(inputs):
|
|
filepath = DATA_DIR / f"{self.name}_input_{i}.bin"
|
|
write_tensor(filepath, inp)
|
|
input_files.append(str(filepath))
|
|
|
|
# Run Foundry
|
|
output_file = DATA_DIR / f"{self.name}_output.bin"
|
|
success, foundry_time = run_foundry(
|
|
self.foundry_operation(),
|
|
input_files,
|
|
str(output_file),
|
|
self.foundry_params()
|
|
)
|
|
|
|
if not success:
|
|
return CompareResult(
|
|
name=self.name,
|
|
passed=False,
|
|
max_diff=float('inf'),
|
|
mean_diff=float('inf'),
|
|
tolerance=self.tolerance,
|
|
input_shape=list(inputs[0].shape),
|
|
time_pytorch_ms=pytorch_time,
|
|
time_foundry_ms=foundry_time,
|
|
error_message="Foundry execution failed"
|
|
)
|
|
|
|
# Read Foundry output
|
|
actual = read_tensor(output_file)
|
|
|
|
# Compare
|
|
passed, max_diff, mean_diff = compare_tensors(
|
|
expected, actual, self.tolerance
|
|
)
|
|
|
|
# Debug output
|
|
if debug:
|
|
write_tensor(DATA_DIR / f"{self.name}_expected.bin", expected)
|
|
with open(DATA_DIR / f"{self.name}_diff.txt", 'w') as f:
|
|
f.write(f"Max diff: {max_diff}\n")
|
|
f.write(f"Mean diff: {mean_diff}\n")
|
|
f.write(f"Tolerance: {self.tolerance}\n")
|
|
f.write(f"Passed: {passed}\n")
|
|
if not passed:
|
|
diff = np.abs(expected - actual)
|
|
idx = np.unravel_index(np.argmax(diff), diff.shape)
|
|
f.write(f"Max diff at index {idx}\n")
|
|
f.write(f"Expected: {expected[idx]}\n")
|
|
f.write(f"Actual: {actual[idx]}\n")
|
|
|
|
result = CompareResult(
|
|
name=self.name,
|
|
passed=passed,
|
|
max_diff=max_diff,
|
|
mean_diff=mean_diff,
|
|
tolerance=self.tolerance,
|
|
input_shape=list(inputs[0].shape),
|
|
time_pytorch_ms=pytorch_time,
|
|
time_foundry_ms=foundry_time,
|
|
error_message=None if passed else f"max_diff {max_diff} > tolerance {self.tolerance}"
|
|
)
|
|
|
|
self.results.append(result)
|
|
return result
|
|
|
|
except Exception as e:
|
|
return CompareResult(
|
|
name=self.name,
|
|
passed=False,
|
|
max_diff=float('inf'),
|
|
mean_diff=float('inf'),
|
|
tolerance=self.tolerance,
|
|
input_shape=[],
|
|
time_pytorch_ms=0.0,
|
|
time_foundry_ms=0.0,
|
|
error_message=str(e)
|
|
)
|
|
|
|
|
|
class TestRunner:
|
|
"""Run multiple comparison tests and report results."""
|
|
|
|
def __init__(self):
|
|
self.tests: List[ComparisonTest] = []
|
|
self.results: List[CompareResult] = []
|
|
|
|
def add_test(self, test: ComparisonTest):
|
|
self.tests.append(test)
|
|
|
|
def run_all(self, debug: bool = False) -> List[CompareResult]:
|
|
"""Run all tests and return results."""
|
|
self.results = []
|
|
|
|
print("=" * 80)
|
|
print("FOUNDRY vs PYTORCH COMPARISON TESTS")
|
|
print("=" * 80)
|
|
print()
|
|
|
|
for test in self.tests:
|
|
result = test.run(debug=debug)
|
|
self.results.append(result)
|
|
|
|
status = "\033[92m[PASS]\033[0m" if result.passed else "\033[91m[FAIL]\033[0m"
|
|
dots = "." * (40 - len(test.name))
|
|
print(f" {status} {test.name} {dots} max_diff={result.max_diff:.6f}, mean_diff={result.mean_diff:.6f}")
|
|
|
|
if not result.passed and result.error_message:
|
|
print(f" {result.error_message}")
|
|
|
|
print()
|
|
print("=" * 80)
|
|
passed = sum(1 for r in self.results if r.passed)
|
|
failed = len(self.results) - passed
|
|
print(f"SUMMARY: {passed}/{len(self.results)} passed, {failed} failed")
|
|
print("=" * 80)
|
|
|
|
return self.results
|
|
|
|
def save_results(self):
|
|
"""Save results to JSON file."""
|
|
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
filepath = RESULTS_DIR / f"{timestamp}.json"
|
|
|
|
data = {
|
|
"timestamp": datetime.now().isoformat(),
|
|
"results": [asdict(r) for r in self.results]
|
|
}
|
|
|
|
with open(filepath, 'w') as f:
|
|
json.dump(data, f, indent=2)
|
|
|
|
print(f"Results saved to {filepath}")
|
|
|
|
|
|
def print_summary(results: List[CompareResult]):
|
|
"""Print a summary of test results."""
|
|
passed = [r for r in results if r.passed]
|
|
failed = [r for r in results if not r.passed]
|
|
|
|
print("\n" + "=" * 60)
|
|
print("VERIFICATION SUMMARY")
|
|
print("=" * 60)
|
|
print(f"Total: {len(results)}")
|
|
print(f"Passed: {len(passed)} ({100*len(passed)/len(results):.1f}%)")
|
|
print(f"Failed: {len(failed)}")
|
|
|
|
if failed:
|
|
print("\nFailed tests:")
|
|
for r in failed:
|
|
print(f" - {r.name}: max_diff={r.max_diff:.6e} (tolerance={r.tolerance:.0e})")
|
|
if r.error_message:
|
|
print(f" {r.error_message}")
|