"""
triton_svd3.py — Fused Triton SVD kernel for batched M×3 matrices.

Target use case: CIFAR-sized images (32×32=1024 pixels, 3 channels)
where cuSOLVER overhead dominates because the "thin" dimension is only 3.

Architecture:
  1. Each program handles one (M×3) matrix from the batch
  2. Compute 3×3 Gram matrix G = A^T A  (6 unique values, in registers)
  3. Diagonalize G via Jacobi rotations   (3×3 converges in ≤6 sweeps)
  4. S = sqrt(eigenvalues), V = eigenvectors
  5. U = A @ V @ diag(1/S)               (tiled reduction over M)

The entire 3×3 eigensolver lives in scalar registers — zero shared memory,
zero global memory round-trips. The only bandwidth cost is loading A and
writing back U, S, Vh.

Author: AbstractPhil / Claude
"""

import triton
import triton.language as tl
import torch
import math


# ============================================================================
# Core kernel: batched SVD for (B, M, 3) tensors
# ============================================================================

@triton.jit
def _svd3_kernel(
    # Pointers
    A_ptr,      # (B, M, 3) input
    U_ptr,      # (B, M, 3) left singular vectors (thin)
    S_ptr,      # (B, 3)    singular values
    Vh_ptr,     # (B, 3, 3) right singular vectors transposed
    # Dimensions
    M: tl.constexpr,          # spatial dim (1024 for CIFAR)
    BLOCK_M: tl.constexpr,    # tile size for M-dimension loads
    JACOBI_ITERS: tl.constexpr,  # number of full Jacobi sweeps
    EPS: tl.constexpr,        # numerical floor
):
    """
    One program instance = one (M, 3) matrix in the batch.
    """
    bid = tl.program_id(0)  # batch index

    # ---------------------------------------------------------------
    # Stage 1: Compute 3×3 Gram matrix G = A^T @ A
    #
    # G is symmetric, so we only need 6 accumulators:
    #   g00 g01 g02
    #       g11 g12
    #           g22
    # ---------------------------------------------------------------
    g00 = tl.zeros([], dtype=tl.float32)
    g01 = tl.zeros([], dtype=tl.float32)
    g02 = tl.zeros([], dtype=tl.float32)
    g11 = tl.zeros([], dtype=tl.float32)
    g12 = tl.zeros([], dtype=tl.float32)
    g22 = tl.zeros([], dtype=tl.float32)

    base = bid * M * 3

    # Tiled accumulation over the spatial dimension
    for block_start in range(0, M, BLOCK_M):
        offs = tl.arange(0, BLOCK_M)  # [0, 1, ..., BLOCK_M-1]
        row_idx = block_start + offs   # actual row indices
        mask = row_idx < M

        # Load 3 columns for this tile
        # A[bid, row, c] = A_ptr[bid*M*3 + row*3 + c]
        ptr0 = base + row_idx * 3 + 0
        ptr1 = base + row_idx * 3 + 1
        ptr2 = base + row_idx * 3 + 2

        a0 = tl.load(A_ptr + ptr0, mask=mask, other=0.0).to(tl.float32)
        a1 = tl.load(A_ptr + ptr1, mask=mask, other=0.0).to(tl.float32)
        a2 = tl.load(A_ptr + ptr2, mask=mask, other=0.0).to(tl.float32)

        # Accumulate outer products
        g00 += tl.sum(a0 * a0)
        g01 += tl.sum(a0 * a1)
        g02 += tl.sum(a0 * a2)
        g11 += tl.sum(a1 * a1)
        g12 += tl.sum(a1 * a2)
        g22 += tl.sum(a2 * a2)

    # ---------------------------------------------------------------
    # Stage 2: 3×3 Jacobi eigendecomposition (all in scalars)
    #
    # Cyclic Jacobi: rotate pairs (0,1), (0,2), (1,2) each sweep.
    # For 3×3 symmetric PSD, 4-6 sweeps is overkill-level convergence.
    #
    # We maintain:
    #   - g_ij: the evolving matrix entries (symmetric, 6 values)
    #   - v_ij: the eigenvector matrix (9 values, starts as I)
    # ---------------------------------------------------------------

    # Eigenvector accumulator V (row-major: v_rc = V[r,c])
    v00 = 1.0; v01 = 0.0; v02 = 0.0
    v10 = 0.0; v11 = 1.0; v12 = 0.0
    v20 = 0.0; v21 = 0.0; v22 = 1.0

    for _sweep in range(JACOBI_ITERS):
        # --- Rotation (p=0, q=1): zero out g01 ---
        #
        # Jacobi rotation angle:
        #   if |g_pq| < eps: skip
        #   tau = (g_qq - g_pp) / (2 * g_pq)
        #   t = sign(tau) / (|tau| + sqrt(1 + tau^2))
        #   c = 1/sqrt(1+t^2),  s = t*c

        # -- pair (0, 1) --
        off_diag = g01
        diag_diff = g11 - g00
        abs_off = tl.abs(off_diag)
        # Compute rotation
        tau_01 = tl.where(abs_off > EPS, diag_diff / (2.0 * off_diag), 0.0)
        abs_tau = tl.abs(tau_01)
        t_01 = tl.where(
            abs_off > EPS,
            tl.where(tau_01 >= 0, 1.0, -1.0) / (abs_tau + tl.sqrt(1.0 + tau_01 * tau_01)),
            0.0,
        )
        c01 = 1.0 / tl.sqrt(1.0 + t_01 * t_01)
        s01 = t_01 * c01

        # Apply Givens to G (symmetric update for p=0, q=1)
        new_g00 = c01*c01*g00 - 2.0*s01*c01*g01 + s01*s01*g11
        new_g11 = s01*s01*g00 + 2.0*s01*c01*g01 + c01*c01*g11
        new_g01 = 0.0  # This is the point
        new_g02 = c01*g02 - s01*g12
        new_g12 = s01*g02 + c01*g12
        # g22 unchanged
        g00 = new_g00; g11 = new_g11; g01 = new_g01
        g02 = new_g02; g12 = new_g12

        # Apply to V columns: V[:, 0], V[:, 1]
        nv00 = c01*v00 - s01*v01;  nv01 = s01*v00 + c01*v01
        nv10 = c01*v10 - s01*v11;  nv11 = s01*v10 + c01*v11
        nv20 = c01*v20 - s01*v21;  nv21 = s01*v20 + c01*v21
        v00 = nv00; v01 = nv01; v10 = nv10; v11 = nv11; v20 = nv20; v21 = nv21

        # -- pair (0, 2) --
        off_diag = g02
        diag_diff = g22 - g00
        abs_off = tl.abs(off_diag)
        tau_02 = tl.where(abs_off > EPS, diag_diff / (2.0 * off_diag), 0.0)
        abs_tau = tl.abs(tau_02)
        t_02 = tl.where(
            abs_off > EPS,
            tl.where(tau_02 >= 0, 1.0, -1.0) / (abs_tau + tl.sqrt(1.0 + tau_02 * tau_02)),
            0.0,
        )
        c02 = 1.0 / tl.sqrt(1.0 + t_02 * t_02)
        s02 = t_02 * c02

        new_g00 = c02*c02*g00 - 2.0*s02*c02*g02 + s02*s02*g22
        new_g22 = s02*s02*g00 + 2.0*s02*c02*g02 + c02*c02*g22
        new_g02 = 0.0
        new_g01 = c02*g01 - s02*g12
        new_g12_b = s02*g01 + c02*g12
        g00 = new_g00; g22 = new_g22; g02 = new_g02
        g01 = new_g01; g12 = new_g12_b

        nv00 = c02*v00 - s02*v02;  nv02 = s02*v00 + c02*v02
        nv10 = c02*v10 - s02*v12;  nv12 = s02*v10 + c02*v12
        nv20 = c02*v20 - s02*v22;  nv22 = s02*v20 + c02*v22
        v00 = nv00; v02 = nv02; v10 = nv10; v12 = nv12; v20 = nv20; v22 = nv22

        # -- pair (1, 2) --
        off_diag = g12
        diag_diff = g22 - g11
        abs_off = tl.abs(off_diag)
        tau_12 = tl.where(abs_off > EPS, diag_diff / (2.0 * off_diag), 0.0)
        abs_tau = tl.abs(tau_12)
        t_12 = tl.where(
            abs_off > EPS,
            tl.where(tau_12 >= 0, 1.0, -1.0) / (abs_tau + tl.sqrt(1.0 + tau_12 * tau_12)),
            0.0,
        )
        c12 = 1.0 / tl.sqrt(1.0 + t_12 * t_12)
        s12 = t_12 * c12

        new_g11 = c12*c12*g11 - 2.0*s12*c12*g12 + s12*s12*g22
        new_g22 = s12*s12*g11 + 2.0*s12*c12*g12 + c12*c12*g22
        new_g12 = 0.0
        new_g01 = c12*g01 - s12*g02
        new_g02_b = s12*g01 + c12*g02
        g11 = new_g11; g22 = new_g22; g12 = new_g12
        g01 = new_g01; g02 = new_g02_b

        nv01 = c12*v01 - s12*v02;  nv02 = s12*v01 + c12*v02
        nv11 = c12*v11 - s12*v12;  nv12 = s12*v11 + c12*v12
        nv21 = c12*v21 - s12*v22;  nv22 = s12*v21 + c12*v22
        v01 = nv01; v02 = nv02; v11 = nv11; v12 = nv12; v21 = nv21; v22 = nv22

    # ---------------------------------------------------------------
    # Stage 2b: Extract eigenvalues, sort descending
    # Diagonal of G now holds eigenvalues of A^T A
    # ---------------------------------------------------------------
    eig0 = tl.maximum(g00, EPS)
    eig1 = tl.maximum(g11, EPS)
    eig2 = tl.maximum(g22, EPS)

    s0 = tl.sqrt(eig0)
    s1 = tl.sqrt(eig1)
    s2 = tl.sqrt(eig2)

    # Sorting network for 3 elements (descending)
    # We need to sort S and permute V columns accordingly.
    # Approach: compare-and-swap on the scalar eigenvalues + V columns
    #
    # 3-element sorting network: (0,1), (0,2), (1,2)

    # swap(0, 1) if s0 < s1
    do_swap = s0 < s1
    s0, s1 = tl.where(do_swap, s1, s0), tl.where(do_swap, s0, s1)
    # Swap V columns 0 and 1
    tv00 = v00; tv10 = v10; tv20 = v20
    v00 = tl.where(do_swap, v01, v00); v01 = tl.where(do_swap, tv00, v01)
    v10 = tl.where(do_swap, v11, v10); v11 = tl.where(do_swap, tv10, v11)
    v20 = tl.where(do_swap, v21, v20); v21 = tl.where(do_swap, tv20, v21)

    # swap(0, 2) if s0 < s2
    do_swap = s0 < s2
    s0, s2 = tl.where(do_swap, s2, s0), tl.where(do_swap, s0, s2)
    tv00 = v00; tv10 = v10; tv20 = v20
    v00 = tl.where(do_swap, v02, v00); v02 = tl.where(do_swap, tv00, v02)
    v10 = tl.where(do_swap, v12, v10); v12 = tl.where(do_swap, tv10, v12)
    v20 = tl.where(do_swap, v22, v20); v22 = tl.where(do_swap, tv20, v22)

    # swap(1, 2) if s1 < s2
    do_swap = s1 < s2
    s1, s2 = tl.where(do_swap, s2, s1), tl.where(do_swap, s1, s2)
    tv01 = v01; tv11 = v11; tv21 = v21
    v01 = tl.where(do_swap, v02, v01); v02 = tl.where(do_swap, tv01, v02)
    v11 = tl.where(do_swap, v12, v11); v12 = tl.where(do_swap, tv11, v12)
    v21 = tl.where(do_swap, v22, v21); v22 = tl.where(do_swap, tv21, v22)

    # ---------------------------------------------------------------
    # Stage 2c: Write S and Vh
    # ---------------------------------------------------------------
    s_base = bid * 3
    tl.store(S_ptr + s_base + 0, s0)
    tl.store(S_ptr + s_base + 1, s1)
    tl.store(S_ptr + s_base + 2, s2)

    # Vh = V^T  (Vh[i,j] = V[j,i])
    vh_base = bid * 9
    tl.store(Vh_ptr + vh_base + 0, v00)  # Vh[0,0] = V[0,0]
    tl.store(Vh_ptr + vh_base + 1, v10)  # Vh[0,1] = V[1,0]
    tl.store(Vh_ptr + vh_base + 2, v20)  # Vh[0,2] = V[2,0]
    tl.store(Vh_ptr + vh_base + 3, v01)  # Vh[1,0] = V[0,1]
    tl.store(Vh_ptr + vh_base + 4, v11)  # Vh[1,1] = V[1,1]
    tl.store(Vh_ptr + vh_base + 5, v21)  # Vh[1,2] = V[2,1]
    tl.store(Vh_ptr + vh_base + 6, v02)  # Vh[2,0] = V[0,2]
    tl.store(Vh_ptr + vh_base + 7, v12)  # Vh[2,1] = V[1,2]
    tl.store(Vh_ptr + vh_base + 8, v22)  # Vh[2,2] = V[2,2]

    # ---------------------------------------------------------------
    # Stage 3: Recover U = A @ V @ diag(1/S)
    #
    # U[:, c] = (1/S[c]) * A @ V[:, c]
    # Tiled over M to keep memory pressure bounded.
    # ---------------------------------------------------------------
    inv_s0 = 1.0 / (s0 + EPS)
    inv_s1 = 1.0 / (s1 + EPS)
    inv_s2 = 1.0 / (s2 + EPS)

    for block_start in range(0, M, BLOCK_M):
        offs = tl.arange(0, BLOCK_M)
        row_idx = block_start + offs
        mask = row_idx < M

        ptr0 = base + row_idx * 3 + 0
        ptr1 = base + row_idx * 3 + 1
        ptr2 = base + row_idx * 3 + 2

        a0 = tl.load(A_ptr + ptr0, mask=mask, other=0.0).to(tl.float32)
        a1 = tl.load(A_ptr + ptr1, mask=mask, other=0.0).to(tl.float32)
        a2 = tl.load(A_ptr + ptr2, mask=mask, other=0.0).to(tl.float32)

        # U[:, 0] = (A @ V[:, 0]) / s0
        u0 = (a0 * v00 + a1 * v10 + a2 * v20) * inv_s0
        # U[:, 1] = (A @ V[:, 1]) / s1
        u1 = (a0 * v01 + a1 * v11 + a2 * v21) * inv_s1
        # U[:, 2] = (A @ V[:, 2]) / s2
        u2 = (a0 * v02 + a1 * v12 + a2 * v22) * inv_s2

        u_base = bid * M * 3
        tl.store(U_ptr + u_base + row_idx * 3 + 0, u0, mask=mask)
        tl.store(U_ptr + u_base + row_idx * 3 + 1, u1, mask=mask)
        tl.store(U_ptr + u_base + row_idx * 3 + 2, u2, mask=mask)


# ============================================================================
# Python wrapper
# ============================================================================

def batched_svd3(
    A: torch.Tensor,
    block_m: int = 128,
    jacobi_iters: int = 6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Batched thin SVD for (B, M, 3) float32 tensors.

    Args:
        A: Input tensor of shape (B, M, 3). M can be anything (1024 for CIFAR).
        block_m: Tile size for the spatial dimension. 128 is good for M=1024.
        jacobi_iters: Number of cyclic Jacobi sweeps. 6 is overkill for 3×3.

    Returns:
        U:  (B, M, 3)  — thin left singular vectors
        S:  (B, 3)     — singular values, descending
        Vh: (B, 3, 3)  — right singular vectors transposed
    """
    assert A.ndim == 3 and A.shape[2] == 3, f"Expected (B, M, 3), got {A.shape}"
    assert A.is_cuda, "Input must be on CUDA"

    B, M, _ = A.shape
    A_f32 = A.contiguous().float()

    U = torch.empty((B, M, 3), dtype=torch.float32, device=A.device)
    S = torch.empty((B, 3), dtype=torch.float32, device=A.device)
    Vh = torch.empty((B, 3, 3), dtype=torch.float32, device=A.device)

    _svd3_kernel[(B,)](
        A_f32, U, S, Vh,
        M=M,
        BLOCK_M=block_m,
        JACOBI_ITERS=jacobi_iters,
        EPS=1e-12,
    )

    return U, S, Vh


# ============================================================================
# Benchmark & validation harness
# ============================================================================

def _test_correctness(B=256, M=1024):
    """Validate against torch.linalg.svd."""
    A = torch.randn(B, M, 3, device="cuda", dtype=torch.float32)

    U, S, Vh = batched_svd3(A)

    # Reference: torch.linalg.svd on 3D is batched
    U_ref, S_ref, Vh_ref = torch.linalg.svd(A, full_matrices=False)

    # Singular values should match (tight — these are just 3 values)
    s_err = (S - S_ref).abs().max().item()
    print(f"[correctness] S max error:     {s_err:.2e}")

    # Reconstruction: A ≈ U @ diag(S) @ Vh
    # Compare against torch's own recon error as the floor —
    # f32 accumulation over M=1024 rows means ~1e-3 is physical.
    recon = torch.bmm(U * S.unsqueeze(1), Vh)
    recon_ref = torch.bmm(U_ref * S_ref.unsqueeze(1), Vh_ref)
    r_err = (A - recon).abs().max().item()
    r_ref_err = (A - recon_ref).abs().max().item()
    print(f"[correctness] Recon error:     {r_err:.2e}  (ref: {r_ref_err:.2e})")

    # Orthogonality: U^T U ≈ I
    UtU = torch.bmm(U.transpose(1, 2), U)
    eye = torch.eye(3, device="cuda").expand(B, -1, -1)
    orth_err = (UtU - eye).abs().max().item()
    print(f"[correctness] U orthog error:  {orth_err:.2e}")

    # S should be tight; recon only needs to match torch's own floor
    assert s_err < 1e-3, f"Singular value error {s_err} > 1e-3"
    assert r_err < r_ref_err * 2.0 + 1e-5, (
        f"Recon error {r_err:.2e} > 2x torch ref {r_ref_err:.2e}"
    )
    assert orth_err < 1e-3, f"Orthogonality error {orth_err} > 1e-3"
    print("[correctness] PASSED\n")


def _cuda_timer(fn, warmup=25, iters=100):
    """
    CUDA-event-timed benchmark. Returns (mean_ms, std_ms) over `iters` runs.
    Uses per-iteration events for proper distribution stats.
    """
    # Warmup — let triton autotuner / cublas handle settle
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()

    starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
    ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]

    for i in range(iters):
        starts[i].record()
        fn()
        ends[i].record()
    torch.cuda.synchronize()

    times = [starts[i].elapsed_time(ends[i]) for i in range(iters)]
    t = torch.tensor(times)
    return t.mean().item(), t.std().item(), t.median().item(), t.min().item(), t.max().item()


def _profile_sweep():
    """
    Full profiling sweep: batch sizes × spatial dims × block sizes.
    Compares Triton SVD3 vs torch.linalg.svd with CUDA event timing.
    """
    import json

    device_name = torch.cuda.get_device_name(0)
    print(f"{'='*72}")
    print(f"  Triton SVD3 Profiling — {device_name}")
    print(f"{'='*72}\n")

    # ------------------------------------------------------------------
    # Sweep 1: Batch scaling at fixed M=1024 (CIFAR spatial)
    # ------------------------------------------------------------------
    print(f"{'─'*72}")
    print(f"  SWEEP 1: Batch scaling  (M=1024, BLOCK_M=128)")
    print(f"{'─'*72}")
    print(f"  {'B':>6}  {'Triton ms':>12} {'±std':>8} {'Torch ms':>12} {'±std':>8} {'Speedup':>8}")
    print(f"  {'─'*62}")

    batch_sizes = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
    M = 1024
    batch_results = []

    for B in batch_sizes:
        A = torch.randn(B, M, 3, device="cuda", dtype=torch.float32)

        tri_mean, tri_std, tri_med, tri_min, tri_max = _cuda_timer(
            lambda: batched_svd3(A, block_m=128)
        )
        tch_mean, tch_std, tch_med, tch_min, tch_max = _cuda_timer(
            lambda: torch.linalg.svd(A, full_matrices=False)
        )
        speedup = tch_mean / (tri_mean + 1e-9)

        print(f"  {B:>6}  {tri_mean:>10.3f}ms {tri_std:>6.3f}  {tch_mean:>10.3f}ms {tch_std:>6.3f}  {speedup:>7.2f}x")
        batch_results.append({
            "B": B, "M": M,
            "triton_mean_ms": round(tri_mean, 4), "triton_std_ms": round(tri_std, 4),
            "triton_median_ms": round(tri_med, 4), "triton_min_ms": round(tri_min, 4),
            "torch_mean_ms": round(tch_mean, 4), "torch_std_ms": round(tch_std, 4),
            "torch_median_ms": round(tch_med, 4), "torch_min_ms": round(tch_min, 4),
            "speedup": round(speedup, 3),
        })
        del A
        torch.cuda.empty_cache()

    # ------------------------------------------------------------------
    # Sweep 2: Spatial dim scaling at fixed B=1024
    # ------------------------------------------------------------------
    print(f"\n{'─'*72}")
    print(f"  SWEEP 2: Spatial scaling  (B=1024, BLOCK_M=128)")
    print(f"{'─'*72}")
    print(f"  {'M':>6}  {'Triton ms':>12} {'±std':>8} {'Torch ms':>12} {'±std':>8} {'Speedup':>8}")
    print(f"  {'─'*62}")

    spatial_dims = [64, 256, 512, 1024, 2048, 4096]  # 8×8 to 64×64
    B = 1024
    spatial_results = []

    for M in spatial_dims:
        A = torch.randn(B, M, 3, device="cuda", dtype=torch.float32)

        tri_mean, tri_std, tri_med, tri_min, tri_max = _cuda_timer(
            lambda: batched_svd3(A, block_m=128)
        )
        tch_mean, tch_std, tch_med, tch_min, tch_max = _cuda_timer(
            lambda: torch.linalg.svd(A, full_matrices=False)
        )
        speedup = tch_mean / (tri_mean + 1e-9)

        equiv_hw = int(M**0.5)
        tag = f"~{equiv_hw}x{equiv_hw}" if equiv_hw * equiv_hw == M else f"  {M}"
        print(f"  {tag:>6}  {tri_mean:>10.3f}ms {tri_std:>6.3f}  {tch_mean:>10.3f}ms {tch_std:>6.3f}  {speedup:>7.2f}x")
        spatial_results.append({
            "B": B, "M": M,
            "triton_mean_ms": round(tri_mean, 4), "torch_mean_ms": round(tch_mean, 4),
            "speedup": round(speedup, 3),
        })
        del A
        torch.cuda.empty_cache()

    # ------------------------------------------------------------------
    # Sweep 3: BLOCK_M tuning at fixed B=4096, M=1024
    # ------------------------------------------------------------------
    print(f"\n{'─'*72}")
    print(f"  SWEEP 3: BLOCK_M tuning  (B=4096, M=1024)")
    print(f"{'─'*72}")
    print(f"  {'BLOCK_M':>8}  {'Triton ms':>12} {'±std':>8} {'tiles/img':>10}")
    print(f"  {'─'*44}")

    block_sizes = [32, 64, 128, 256, 512, 1024]
    B, M = 4096, 1024
    A = torch.randn(B, M, 3, device="cuda", dtype=torch.float32)
    block_results = []

    for bm in block_sizes:
        tri_mean, tri_std, tri_med, tri_min, tri_max = _cuda_timer(
            lambda: batched_svd3(A, block_m=bm)
        )
        n_tiles = (M + bm - 1) // bm
        print(f"  {bm:>8}  {tri_mean:>10.3f}ms {tri_std:>6.3f}  {n_tiles:>10}")
        block_results.append({
            "block_m": bm, "triton_mean_ms": round(tri_mean, 4),
            "triton_std_ms": round(tri_std, 4), "n_tiles": n_tiles,
        })

    del A
    torch.cuda.empty_cache()

    # ------------------------------------------------------------------
    # Sweep 4: Throughput — images/sec at peak batch
    # ------------------------------------------------------------------
    print(f"\n{'─'*72}")
    print(f"  SWEEP 4: Throughput (images/sec)")
    print(f"{'─'*72}")

    for B in [4096, 16384]:
        A = torch.randn(B, 1024, 3, device="cuda", dtype=torch.float32)
        tri_mean, *_ = _cuda_timer(lambda: batched_svd3(A, block_m=128))
        tch_mean, *_ = _cuda_timer(lambda: torch.linalg.svd(A, full_matrices=False))

        tri_ips = B / (tri_mean / 1000)
        tch_ips = B / (tch_mean / 1000)
        print(f"  B={B:>5}: Triton {tri_ips:>12,.0f} img/s | Torch {tch_ips:>12,.0f} img/s")
        del A
        torch.cuda.empty_cache()

    # ------------------------------------------------------------------
    # Memory: peak allocation comparison
    # ------------------------------------------------------------------
    print(f"\n{'─'*72}")
    print(f"  MEMORY: Peak allocation (B=4096, M=1024)")
    print(f"{'─'*72}")

    B, M = 4096, 1024
    A = torch.randn(B, M, 3, device="cuda", dtype=torch.float32)

    torch.cuda.reset_peak_memory_stats()
    _ = batched_svd3(A)
    torch.cuda.synchronize()
    tri_peak = torch.cuda.max_memory_allocated() / 1024**2

    torch.cuda.reset_peak_memory_stats()
    _ = torch.linalg.svd(A, full_matrices=False)
    torch.cuda.synchronize()
    tch_peak = torch.cuda.max_memory_allocated() / 1024**2

    print(f"  Triton: {tri_peak:.1f} MB")
    print(f"  Torch:  {tch_peak:.1f} MB")
    print(f"  Ratio:  {tch_peak / (tri_peak + 1e-9):.2f}x\n")

    # ------------------------------------------------------------------
    # Dump JSON for further analysis
    # ------------------------------------------------------------------
    report = {
        "device": device_name,
        "batch_sweep": batch_results,
        "spatial_sweep": spatial_results,
        "block_m_sweep": block_results,
    }
    with open("svd3_profile.json", "w") as f:
        json.dump(report, f, indent=2)
    print(f"  Full results written to svd3_profile.json\n")


if __name__ == "__main__":
    _test_correctness()
    _profile_sweep()