Make svd_pullback! and eigh_pullback! cost proportional to the number of cotangent columns - #292

# Views vs copies of the kept columns in svd_pullback!/eigh_pullback! (untracked, for the review reply):
#   1. mul! with a view of the columns ind as one factor, for ind a range and a Vector
#   2. the pullbacks of this branch (views of U, Vᴴ and V) against the same code with copies
#      `U[:, ind′]`, `Vᴴ[ind′, :]`, `V[:, ind′]` (module Copies), timed alternately in one process
using LinearAlgebra, Random, Printf
using MatrixAlgebraKit
using MatrixAlgebraKit: svd_pullback!, eigh_pullback!
BLAS.set_num_threads(4)
module Copies
using LinearAlgebra, MatrixAlgebraKit
const SUBS = ("view(U, :, ind′)" => "U[:, ind′]", "view(Vᴴ, ind′, :)" => "Vᴴ[ind′, :]", "view(V, :, ind′)" => "V[:, ind′]")
const SRC = map(("svd", "eigh")) do f
    src = read(joinpath(pkgdir(MatrixAlgebraKit), "src", "pullbacks", "$f.jl"), String)
    any(occursin(first(s), src) for s in SUBS) || error("no view in $f.jl")
    return replace(src, SUBS...)
end
const OWN = Set(Symbol(something(m[1], m[2])) for src in SRC for m in eachmatch(r"^\s*function ([^\s(]+)\(|^(?:function )?([A-Za-z_][^\s(]*)\([^\n]*\)\s*=[^=]"m, src))
for name in names(MatrixAlgebraKit; all = true)
    (name in OWN || startswith(string(name), "#") || isdefined(@__MODULE__, name)) && continue
    @eval const $name = MatrixAlgebraKit.$name
end
foreach(src -> include_string(@__MODULE__, src), SRC)
end
function mintimes(f, g; mintotal = 2.0, minreps = 5, maxreps = 15)
    f(); g(); tf = tg = Inf; tot = 0.0
    for i in 1:maxreps
        s = time_ns(); f(); df = (time_ns() - s) / 1.0e9
        s = time_ns(); g(); dg = (time_ns() - s) / 1.0e9
        tf, tg = min(tf, df), min(tg, dg); tot += df + dg
        i >= minreps && tot > mintotal && break
    end
    return tf, tg
end
noimagdiag!(ΔX, X) = (ΔX .-= X .* transpose(im .* imag.(diag(X' * ΔX))); ΔX)
noimagdiag!(ΔX::AbstractMatrix{<:Real}, X) = ΔX
indkinds(n, k, rng) = (("range", 1:k), ("vector", collect(1:k)), ("perm", randperm(rng, n)[1:k]))

println("== mul!(C, A, B') with B = view(V, :, ind) or V[:, ind], n = 1000, BLAS threads 4")
let rng = Xoshiro(1), n = 1000
    for T in (Float64, ComplexF64), k in (50, 200)
        A = randn(rng, T, n, k); V = randn(rng, T, n, n); C = zeros(T, n, n)
        for (name, ind) in indkinds(n, k, rng)
            tv, tc = mintimes(() -> mul!(C, A, view(V, :, ind)', 1, 1), () -> mul!(C, A, V[:, ind]', 1, 1))
            @printf("MUL T=%s k=%d ind=%s view=%.3e copy=%.3e view/copy=%.1f\n", T, k, name, tv, tc, tv / tc)
        end
    end
end

println("== pullbacks: views (branch) vs copies, svd m=1200 n=1000, eigh n=1000")
let rng = Xoshiro(2), m = 1200, n = 1000
    for T in (Float64, ComplexF64)
        A = randn(rng, T, m, n); U, S, Vᴴ = svd_compact(A); Sd = diag(S)
        H = randn(rng, T, n, n); H = H + H'; D, V = eigh_full(H); Dd = diag(D)
        for k in (50, 200, 700), (name, ind) in indkinds(n, k, rng)
            ΔU = noimagdiag!(randn(rng, T, m, k), U[:, ind])
            ΔVᴴ = copy(noimagdiag!(randn(rng, T, n, k), Vᴴ[ind, :]')')
            ΔS = Diagonal(randn(rng, real(T), k))
            f(pb) = () -> pb(zero(A), A, (U, S, Vᴴ), (ΔU, ΔS, ΔVᴴ), ind)
            err = norm(f(svd_pullback!)() - f(Copies.svd_pullback!)())
            tv, tc = mintimes(f(svd_pullback!), f(Copies.svd_pullback!))
            @printf("PB dec=svd T=%s k=%d ind=%s views=%.3e copies=%.3e views/copies=%.2f diff=%.0e\n", T, k, name, tv, tc, tv / tc, err)
            ΔV = noimagdiag!(randn(rng, T, n, k), V[:, ind]); ΔD = Diagonal(randn(rng, real(T), k))
            g(pb) = () -> pb(zero(H), H, (D, V), (ΔD, ΔV), ind)
            err = norm(g(eigh_pullback!)() - g(Copies.eigh_pullback!)())
            tv, tc = mintimes(g(eigh_pullback!), g(Copies.eigh_pullback!))
            @printf("PB dec=eigh T=%s k=%d ind=%s views=%.3e copies=%.3e views/copies=%.2f diff=%.0e\n", T, k, name, tv, tc, tv / tc, err)
        end
    end
end