test_fft.py

"""
Test script comparing FFT implementations:
- NumPy (CPU reference)
- CuPy (GPU)
- PyTorch (GPU)
- cuTile (GPU, from FFT.py)

Tests 1D, 2D FFT at various sizes with timing.
"""

import numpy as np
import cupy as cp
import torch
import time
from FFT import cutile_fft


def sync_and_time():
    """Synchronize CUDA and return current time."""
    torch.cuda.synchronize()
    cp.cuda.Stream.null.synchronize()
    return time.perf_counter()


def cutile_fft_1d_chunked(x: torch.Tensor, factors: tuple, atom_packing_dim: int, chunk_size: int = 8) -> torch.Tensor:
    """
    Perform batched 1D FFT using cuTile, processing in chunks to handle large batch sizes.
    
    The cuTile kernel uses batch size as a compile-time constant for tile shapes,
    so large batch sizes can cause issues. This wrapper chunks the input.
    
    Args:
        x: Input tensor of shape (Batch, N) with complex64 dtype on CUDA
        factors: FFT factorization tuple
        atom_packing_dim: Packing dimension for cuTile kernel
        chunk_size: Maximum batch size per kernel call
    
    Returns:
        FFT result of shape (Batch, N)
    """
    batch_size, N = x.shape
    
    if batch_size <= chunk_size:
        return cutile_fft(x, factors=factors, atom_packing_dim=atom_packing_dim)
    
    # Process in chunks
    outputs = []
    for i in range(0, batch_size, chunk_size):
        chunk = x[i:i+chunk_size]
        out_chunk = cutile_fft(chunk, factors=factors, atom_packing_dim=atom_packing_dim)
        outputs.append(out_chunk)
    
    return torch.cat(outputs, dim=0)


def cutile_fft_2d(x: torch.Tensor, factors: tuple, atom_packing_dim: int, chunk_size: int = 8) -> torch.Tensor:
    """
    Perform 2D FFT using cuTile's batched 1D FFT.
    
    Strategy: Apply 1D FFT along rows, then along columns.
    For an (M, N) input:
      1. FFT along axis=-1 (rows): treat as M batches of N elements
      2. FFT along axis=-2 (cols): transpose to (N, M), treat as N batches of M elements, transpose back
    
    Args:
        x: Input tensor of shape (M, N) with complex64 dtype on CUDA
        factors: Tuple (F0, F1, F2) for FFT factorization (must work for both M and N)
        atom_packing_dim: Packing dimension for cuTile kernel
        chunk_size: Maximum batch size per kernel call (default 8)
    
    Returns:
        2D FFT result of shape (M, N)
    """
    M, N = x.shape
    
    # Step 1: FFT along rows (axis=-1)
    # Input shape (M, N) is already in batched format for 1D FFT
    x_rows = cutile_fft_1d_chunked(x, factors=factors, atom_packing_dim=atom_packing_dim, chunk_size=chunk_size)
    
    # Step 2: FFT along columns (axis=-2)
    # Transpose to (N, M), apply FFT, transpose back
    x_transposed = x_rows.T.contiguous()  # Shape: (N, M)
    x_cols = cutile_fft_1d_chunked(x_transposed, factors=factors, atom_packing_dim=atom_packing_dim, chunk_size=chunk_size)
    result = x_cols.T.contiguous()  # Shape: (M, N)
    
    return result


def test_1d_fft():
    """Test 1D FFT with 8-element vector."""
    # Configuration for 8-element FFT
    N = 8
    BATCH_SIZE = 1
    FFT_FACTORS = (2, 2, 2)  # 2 * 2 * 2 = 8
    ATOM_PACKING_DIM = 2     # N*2=16, 16 % 2 = 0

    # Set seeds for reproducibility
    np.random.seed(42)
    torch.manual_seed(42)
    cp.random.seed(42)

    # Create input data (complex64)
    # Start with numpy as the ground truth
    input_np = (np.random.randn(BATCH_SIZE, N) + 1j * np.random.randn(BATCH_SIZE, N)).astype(np.complex64)

    print("=" * 60)
    print("TEST 1: 1D FFT - 8 Element Vector")
    print("=" * 60)
    print(f"\nInput shape: {input_np.shape}, dtype: {input_np.dtype}")
    print(f"Input data:\n{input_np}")

    # 1. NumPy FFT (CPU reference)
    print("\n" + "-" * 40)
    print("1. NumPy FFT (CPU reference)")
    output_numpy = np.fft.fft(input_np, axis=-1)
    print(f"Output:\n{output_numpy}")

    # 2. CuPy FFT (GPU)
    print("\n" + "-" * 40)
    print("2. CuPy FFT (GPU)")
    input_cupy = cp.asarray(input_np)
    output_cupy = cp.fft.fft(input_cupy, axis=-1)
    output_cupy_np = cp.asnumpy(output_cupy)
    print(f"Output:\n{output_cupy_np}")

    # Compare with NumPy
    cupy_match = np.allclose(output_numpy, output_cupy_np, rtol=1e-5, atol=1e-5)
    print(f"Matches NumPy: {cupy_match}")

    # 3. PyTorch FFT (GPU)
    print("\n" + "-" * 40)
    print("3. PyTorch FFT (GPU)")
    input_torch = torch.from_numpy(input_np).to('cuda')
    output_torch = torch.fft.fft(input_torch, dim=-1)
    output_torch_np = output_torch.cpu().numpy()
    print(f"Output:\n{output_torch_np}")

    # Compare with NumPy
    torch_match = np.allclose(output_numpy, output_torch_np, rtol=1e-5, atol=1e-5)
    print(f"Matches NumPy: {torch_match}")

    # 4. cuTile FFT (GPU)
    print("\n" + "-" * 40)
    print("4. cuTile FFT (GPU)")
    # cuTile expects torch tensor on CUDA
    input_cutile = torch.from_numpy(input_np).to('cuda')
    output_cutile = cutile_fft(
        x=input_cutile,
        factors=FFT_FACTORS,
        atom_packing_dim=ATOM_PACKING_DIM
    )
    output_cutile_np = output_cutile.cpu().numpy()
    print(f"Output:\n{output_cutile_np}")

    # Compare with NumPy
    cutile_match = np.allclose(output_numpy, output_cutile_np, rtol=1e-5, atol=1e-5)
    print(f"Matches NumPy: {cutile_match}")

    # Summary
    print("\n" + "-" * 40)
    print("1D FFT SUMMARY")
    print("-" * 40)
    print(f"CuPy    matches NumPy: {cupy_match}")
    print(f"PyTorch matches NumPy: {torch_match}")
    print(f"cuTile  matches NumPy: {cutile_match}")

    all_match = cupy_match and torch_match and cutile_match
    if all_match:
        print("\nAll 1D implementations match!")
    else:
        print("\nSome implementations differ - check outputs above.")
        print("\nMax absolute differences from NumPy:")
        print(f"  CuPy:    {np.max(np.abs(output_numpy - output_cupy_np))}")
        print(f"  PyTorch: {np.max(np.abs(output_numpy - output_torch_np))}")
        print(f"  cuTile:  {np.max(np.abs(output_numpy - output_cutile_np))}")

    return all_match


def test_2d_fft():
    """Test 2D FFT with 8x8 matrix."""
    # Configuration for 8x8 2D FFT
    M, N = 8, 8
    FFT_FACTORS = (2, 2, 2)  # 2 * 2 * 2 = 8 (works for both dimensions)
    ATOM_PACKING_DIM = 2     # N*2=16, 16 % 2 = 0

    # Set seeds for reproducibility
    np.random.seed(123)
    torch.manual_seed(123)
    cp.random.seed(123)

    # Create input data (complex64)
    input_np = (np.random.randn(M, N) + 1j * np.random.randn(M, N)).astype(np.complex64)

    print("\n" + "=" * 60)
    print("TEST 2: 2D FFT - 8x8 Matrix")
    print("=" * 60)
    print(f"\nInput shape: {input_np.shape}, dtype: {input_np.dtype}")
    print(f"Input data (first 2 rows):\n{input_np[:2, :]}")

    # 1. NumPy 2D FFT (CPU reference)
    print("\n" + "-" * 40)
    print("1. NumPy FFT2 (CPU reference)")
    output_numpy = np.fft.fft2(input_np)
    print(f"Output (first 2 rows):\n{output_numpy[:2, :]}")

    # 2. CuPy 2D FFT (GPU)
    print("\n" + "-" * 40)
    print("2. CuPy FFT2 (GPU)")
    input_cupy = cp.asarray(input_np)
    output_cupy = cp.fft.fft2(input_cupy)
    output_cupy_np = cp.asnumpy(output_cupy)
    print(f"Output (first 2 rows):\n{output_cupy_np[:2, :]}")

    cupy_match = np.allclose(output_numpy, output_cupy_np, rtol=1e-5, atol=1e-5)
    print(f"Matches NumPy: {cupy_match}")

    # 3. PyTorch 2D FFT (GPU)
    print("\n" + "-" * 40)
    print("3. PyTorch FFT2 (GPU)")
    input_torch = torch.from_numpy(input_np).to('cuda')
    output_torch = torch.fft.fft2(input_torch)
    output_torch_np = output_torch.cpu().numpy()
    print(f"Output (first 2 rows):\n{output_torch_np[:2, :]}")

    torch_match = np.allclose(output_numpy, output_torch_np, rtol=1e-5, atol=1e-5)
    print(f"Matches NumPy: {torch_match}")

    # 4. cuTile 2D FFT (GPU) - using our wrapper
    print("\n" + "-" * 40)
    print("4. cuTile FFT2 (GPU) - via 1D FFT composition")
    input_cutile = torch.from_numpy(input_np).to('cuda')
    output_cutile = cutile_fft_2d(
        x=input_cutile,
        factors=FFT_FACTORS,
        atom_packing_dim=ATOM_PACKING_DIM
    )
    output_cutile_np = output_cutile.cpu().numpy()
    print(f"Output (first 2 rows):\n{output_cutile_np[:2, :]}")

    cutile_match = np.allclose(output_numpy, output_cutile_np, rtol=1e-5, atol=1e-5)
    print(f"Matches NumPy: {cutile_match}")

    # Summary
    print("\n" + "-" * 40)
    print("2D FFT SUMMARY")
    print("-" * 40)
    print(f"CuPy    matches NumPy: {cupy_match}")
    print(f"PyTorch matches NumPy: {torch_match}")
    print(f"cuTile  matches NumPy: {cutile_match}")

    all_match = cupy_match and torch_match and cutile_match
    if all_match:
        print("\nAll 2D implementations match!")
    else:
        print("\nSome implementations differ.")
        print("\nMax absolute differences from NumPy:")
        print(f"  CuPy:    {np.max(np.abs(output_numpy - output_cupy_np))}")
        print(f"  PyTorch: {np.max(np.abs(output_numpy - output_torch_np))}")
        print(f"  cuTile:  {np.max(np.abs(output_numpy - output_cutile_np))}")

    return all_match


def test_2d_fft_timed(size: int, factors: tuple, atom_packing_dim: int, num_warmup: int = 5, num_runs: int = 10):
    """
    Test 2D FFT with timing at specified size.

    Args:
        size: Matrix dimension (size x size)
        factors: FFT factorization tuple (F0, F1, F2) where F0*F1*F2 = size
        atom_packing_dim: Packing dimension for cuTile
        num_warmup: Number of warmup iterations
        num_runs: Number of timed iterations
    """
    M, N = size, size

    # Set seeds for reproducibility
    np.random.seed(456)
    torch.manual_seed(456)
    cp.random.seed(456)

    # Create input data (complex64)
    input_np = (np.random.randn(M, N) + 1j * np.random.randn(M, N)).astype(np.complex64)

    print(f"\n{'=' * 60}")
    print(f"2D FFT BENCHMARK - {size}x{size}")
    print(f"{'=' * 60}")
    print(f"Factors: {factors}, Atom Packing Dim: {atom_packing_dim}")
    print(f"Warmup: {num_warmup}, Timed runs: {num_runs}")

    # Prepare inputs for each implementation
    input_cupy = cp.asarray(input_np)
    input_torch = torch.from_numpy(input_np).to('cuda')
    input_cutile = torch.from_numpy(input_np).to('cuda')

    # ==================== NumPy FFT2 ====================
    print(f"\n{'-' * 40}")
    print("1. NumPy FFT2 (CPU)")

    # Warmup
    for _ in range(num_warmup):
        _ = np.fft.fft2(input_np)

    # Timed runs
    start = time.perf_counter()
    for _ in range(num_runs):
        output_numpy = np.fft.fft2(input_np)
    end = time.perf_counter()
    numpy_time = (end - start) / num_runs * 1000  # ms
    print(f"  Time: {numpy_time:.3f} ms")

    # ==================== CuPy FFT2 ====================
    print(f"\n{'-' * 40}")
    print("2. CuPy FFT2 (GPU)")

    # Warmup
    for _ in range(num_warmup):
        _ = cp.fft.fft2(input_cupy)
    sync_and_time()

    # Timed runs
    start = sync_and_time()
    for _ in range(num_runs):
        output_cupy = cp.fft.fft2(input_cupy)
        cp.cuda.Stream.null.synchronize()  # Sync after each FFT
    end = sync_and_time()
    cupy_time = (end - start) / num_runs * 1000  # ms
    output_cupy_np = cp.asnumpy(output_cupy)
    print(f"  Time: {cupy_time:.3f} ms")

    cupy_match = np.allclose(output_numpy, output_cupy_np, rtol=1e-3, atol=1e-2)
    cupy_max_diff = np.max(np.abs(output_numpy - output_cupy_np))
    print(f"  Matches NumPy: {cupy_match}, max_diff: {cupy_max_diff:.6f}")

    # ==================== PyTorch FFT2 ====================
    print(f"\n{'-' * 40}")
    print("3. PyTorch FFT2 (GPU)")

    # Warmup
    for _ in range(num_warmup):
        _ = torch.fft.fft2(input_torch)
    sync_and_time()

    # Timed runs
    start = sync_and_time()
    for _ in range(num_runs):
        output_torch = torch.fft.fft2(input_torch)
        torch.cuda.synchronize()  # Sync after each FFT
    end = sync_and_time()
    torch_time = (end - start) / num_runs * 1000  # ms
    output_torch_np = output_torch.cpu().numpy()
    print(f"  Time: {torch_time:.3f} ms")

    torch_match = np.allclose(output_numpy, output_torch_np, rtol=1e-3, atol=1e-2)
    torch_max_diff = np.max(np.abs(output_numpy - output_torch_np))
    print(f"  Matches NumPy: {torch_match}, max_diff: {torch_max_diff:.6f}")

    # ==================== cuTile FFT2 ====================
    print(f"\n{'-' * 40}")
    print("4. cuTile FFT2 (GPU)")

    # Warmup
    for _ in range(num_warmup):
        _ = cutile_fft_2d(input_cutile, factors=factors, atom_packing_dim=atom_packing_dim)
    sync_and_time()

    # Timed runs
    start = sync_and_time()
    for _ in range(num_runs):
        output_cutile = cutile_fft_2d(input_cutile, factors=factors, atom_packing_dim=atom_packing_dim)
        torch.cuda.synchronize()  # Sync after each FFT
    end = sync_and_time()
    cutile_time = (end - start) / num_runs * 1000  # ms
    output_cutile_np = output_cutile.cpu().numpy()
    print(f"  Time: {cutile_time:.3f} ms")

    # Use looser tolerance for large FFTs due to float32 accumulation errors
    cutile_match = np.allclose(output_numpy, output_cutile_np, rtol=1e-3, atol=1e-2)
    max_diff = np.max(np.abs(output_numpy - output_cutile_np))
    print(f"  Matches NumPy (rtol=1e-3): {cutile_match}, max_diff: {max_diff:.6f}")

    # ==================== Summary ====================
    print(f"\n{'-' * 40}")
    print(f"TIMING SUMMARY - {size}x{size}")
    print(f"{'-' * 40}")
    print(f"{'Implementation':<20} {'Time (ms)':<12} {'Speedup vs NumPy':<18} {'Correct'}")
    print(f"{'-' * 60}")
    print(f"{'NumPy (CPU)':<20} {numpy_time:<12.3f} {'1.00x':<18} {'--'}")
    print(f"{'CuPy (GPU)':<20} {cupy_time:<12.3f} {numpy_time/cupy_time:<18.2f}x {cupy_match}")
    print(f"{'PyTorch (GPU)':<20} {torch_time:<12.3f} {numpy_time/torch_time:<18.2f}x {torch_match}")
    print(f"{'cuTile (GPU)':<20} {cutile_time:<12.3f} {numpy_time/cutile_time:<18.2f}x {cutile_match}")

    all_match = cupy_match and torch_match and cutile_match

    return {
        'size': size,
        'numpy_time': numpy_time,
        'cupy_time': cupy_time,
        'torch_time': torch_time,
        'cutile_time': cutile_time,
        'all_correct': all_match
    }


def main():
    """Run all FFT tests."""
    print("\n" + "#" * 60)
    print("#" + " " * 20 + "FFT TEST SUITE" + " " * 24 + "#")
    print("#" * 60)

    results = {}

    # Test 1D FFT (8 elements)
    results['1D-8'] = test_1d_fft()

    # Test 2D FFT (8x8)
    results['2D-8x8'] = test_2d_fft()

    # Test 2D FFT with timing at larger sizes
    benchmark_results = []

    # 1024x1024: 1024 = 8 * 8 * 16
    bench_1k = test_2d_fft_timed(
        size=1024,
        factors=(8, 8, 16),
        atom_packing_dim=64,
        num_warmup=3,
        num_runs=10
    )
    benchmark_results.append(bench_1k)
    results['2D-1024x1024'] = bench_1k['all_correct']

    # 2048x2048: 2048 = 8 * 16 * 16
    bench_2k = test_2d_fft_timed(
        size=2048,
        factors=(8, 16, 16),
        atom_packing_dim=64,
        num_warmup=3,
        num_runs=10
    )
    benchmark_results.append(bench_2k)
    results['2D-2048x2048'] = bench_2k['all_correct']

    # Final summary
    print("\n" + "=" * 60)
    print("FINAL SUMMARY")
    print("=" * 60)

    print("\nCorrectness:")
    for test_name, passed in results.items():
        status = "PASSED" if passed else "FAILED"
        print(f"  {test_name}: {status}")

    print("\nPerformance Summary (2D FFT):")
    print(f"{'Size':<12} {'NumPy':<10} {'CuPy':<10} {'PyTorch':<10} {'cuTile':<10}")
    print(f"{'-' * 52}")
    for bench in benchmark_results:
        print(f"{bench['size']}x{bench['size']:<6} "
              f"{bench['numpy_time']:<10.2f} "
              f"{bench['cupy_time']:<10.2f} "
              f"{bench['torch_time']:<10.2f} "
              f"{bench['cutile_time']:<10.2f}")

    all_passed = all(results.values())
    if all_passed:
        print("\nAll tests passed!")
    else:
        print("\nSome tests failed - see details above.")

    return all_passed


if __name__ == "__main__":
    main()