diff --git a/bench-comparison/Main.hs b/bench-comparison/Main.hs new file mode 100644 --- /dev/null +++ b/bench-comparison/Main.hs @@ -0,0 +1,340 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | Cross-library benchmark: linear-massiv vs hmatrix vs linear. +-- +-- Compares performance of numerical linear algebra operations across +-- three Haskell libraries: +-- +-- * __linear-massiv__ — pure Haskell, massiv-backed, type-safe dimensions +-- * __hmatrix__ — FFI to BLAS\/LAPACK (OpenBLAS on this system) +-- * __linear__ — pure Haskell, optimised for small fixed-size (V2–V4) +-- +-- Run with @+RTS -N1@ for fair single-threaded comparison. +module Main (main) where + +import Criterion.Main +import qualified Data.Massiv.Array as M +import GHC.TypeNats (KnownNat) + +-- linear-massiv +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level1 (dotP) +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvecP) +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMulP, matMulPPar) +import Numeric.LinearAlgebra.Massiv.Solve.LU (luSolve, luSolveP) +import Numeric.LinearAlgebra.Massiv.Solve.Cholesky (choleskySolve, choleskySolveP) +import Numeric.LinearAlgebra.Massiv.Orthogonal.QR (qr, qrP) +import Numeric.LinearAlgebra.Massiv.Eigen.Symmetric (symmetricEigen, symmetricEigenP, symmetricEigenPPar, symmetricEigenPDC, tridiagonalizeP) +import Numeric.LinearAlgebra.Massiv.Eigen.SVD (svd, svdP, svdGKP) + +-- hmatrix +import qualified Numeric.LinearAlgebra as H + +-- linear (small fixed-size) +import Linear.V4 (V4(..)) +import qualified Linear.Matrix as LM +import qualified Linear.Metric as LMet + +------------------------------------------------------------------------ +-- linear-massiv matrix generators +------------------------------------------------------------------------ + +mkMatLM :: forall m n. (KnownNat m, KnownNat n) => Matrix m n M.P Double +mkMatLM = makeMatrix @m @n @M.P $ \i j -> + fromIntegral (i * 7 + j * 3 + 1) / 100.0 + +mkVecLM :: forall n. KnownNat n => Vector n M.P Double +mkVecLM = makeVector @n @M.P $ \i -> fromIntegral (i + 1) / 10.0 + +-- Diagonally dominant for LU +mkDDLM :: forall n. KnownNat n => Matrix n n M.P Double +mkDDLM = makeMatrix @n @n @M.P $ \i j -> + fromIntegral (i * 7 + j * 3 + 1) / 100.0 + + if i == j then fromIntegral (dimVal @n) else 0 + +-- SPD: B^T B + nI +mkSPDLM :: forall n. KnownNat n => Matrix n n M.P Double +mkSPDLM = + let nn = dimVal @n + b = makeMatrix @n @n @M.P $ \i j -> + fromIntegral (i * nn + j + 1) / fromIntegral (nn * nn) + in makeMatrix @n @n @M.P $ \i j -> + foldl' (\acc k -> acc + (b ! (i, k)) * (b ! (j, k))) + (if i == j then 1 else 0) + [0..nn-1] + +------------------------------------------------------------------------ +-- hmatrix matrix generators (same numerical entries) +------------------------------------------------------------------------ + +mkMatHM :: Int -> H.Matrix Double +mkMatHM n = (n H.>< n) [ fromIntegral (i * 7 + j * 3 + 1) / 100.0 + | i <- [0..n-1], j <- [0..n-1] ] + +mkVecHM :: Int -> H.Vector Double +mkVecHM n = H.fromList [ fromIntegral (i + 1) / 10.0 | i <- [0..n-1] ] + +mkDDHM :: Int -> H.Matrix Double +mkDDHM n = (n H.>< n) [ fromIntegral (i * 7 + j * 3 + 1) / 100.0 + + if i == j then fromIntegral n else 0 + | i <- [0..n-1], j <- [0..n-1] ] + +mkSPDHM :: Int -> H.Matrix Double +mkSPDHM n = + let b = (n H.>< n) [ fromIntegral (i * n + j + 1) / fromIntegral (n * n) + | i <- [0..n-1], j <- [0..n-1] ] + in H.tr b H.<> b + H.scale (fromIntegral n) (H.ident n) + +------------------------------------------------------------------------ +-- linear (V4) data +------------------------------------------------------------------------ + +linM44a :: V4 (V4 Double) +linM44a = V4 (V4 0.01 0.04 0.07 0.10) + (V4 0.08 0.11 0.14 0.17) + (V4 0.15 0.18 0.21 0.24) + (V4 0.22 0.25 0.28 0.31) + +linM44b :: V4 (V4 Double) +linM44b = V4 (V4 0.34 0.37 0.40 0.43) + (V4 0.41 0.44 0.47 0.50) + (V4 0.48 0.51 0.54 0.57) + (V4 0.55 0.58 0.61 0.64) + +linV4a :: V4 Double +linV4a = V4 0.1 0.2 0.3 0.4 + +linV4b :: V4 Double +linV4b = V4 0.5 0.6 0.7 0.8 + +------------------------------------------------------------------------ +-- hmatrix helpers (avoid operator section issues) +------------------------------------------------------------------------ + +hmGemm :: H.Matrix Double -> H.Matrix Double -> H.Matrix Double +hmGemm = (H.<>) + +hmMatvec :: H.Matrix Double -> H.Vector Double -> H.Vector Double +hmMatvec = (H.#>) + +hmDot :: H.Vector Double -> H.Vector Double -> Double +hmDot = H.dot + +hmLinearSolve :: H.Matrix Double -> H.Vector Double -> H.Matrix Double +hmLinearSolve a b = case H.linearSolve a (H.asColumn b) of + Just x -> x + Nothing -> error "hmLinearSolve: singular matrix" + +hmCholSolve :: H.Matrix Double -> H.Vector Double -> H.Matrix Double +hmCholSolve a b = + let r = H.chol (H.trustSym a) + in H.cholSolve r (H.asColumn b) + +------------------------------------------------------------------------ +-- Benchmarks +------------------------------------------------------------------------ + +main :: IO () +main = do + -- Pre-compute all matrices to avoid construction overhead in benchmarks + let hm4 = mkMatHM 4; hm10 = mkMatHM 10; hm50 = mkMatHM 50 + hm100 = mkMatHM 100; hm200 = mkMatHM 200; hm500 = mkMatHM 500 + hv4 = mkVecHM 4; hv10 = mkVecHM 10; hv50 = mkVecHM 50 + hv100 = mkVecHM 100; hv1000 = mkVecHM 1000 + dd10 = mkDDHM 10; dd50 = mkDDHM 50; dd100 = mkDDHM 100 + spd10 = mkSPDHM 10; spd50 = mkSPDHM 50; spd100 = mkSPDHM 100 + spd200 = mkSPDHM 200; spd500 = mkSPDHM 500 + + defaultMain + [ bgroup "GEMM" + [ bgroup "4x4" + [ bench "linear" $ nf (linM44a LM.!*!) linM44b + , bench "hmatrix" $ nf (hmGemm hm4) hm4 + , bench "linear-massiv" $ nf (matMulP (mkMatLM @4 @4)) (mkMatLM @4 @4) + ] + , bgroup "10x10" + [ bench "hmatrix" $ nf (hmGemm hm10) hm10 + , bench "linear-massiv" $ nf (matMulP (mkMatLM @10 @10)) (mkMatLM @10 @10) + ] + , bgroup "50x50" + [ bench "hmatrix" $ nf (hmGemm hm50) hm50 + , bench "linear-massiv" $ nf (matMulP (mkMatLM @50 @50)) (mkMatLM @50 @50) + ] + , bgroup "100x100" + [ bench "hmatrix" $ nf (hmGemm hm100) hm100 + , bench "linear-massiv" $ nf (matMulP (mkMatLM @100 @100)) (mkMatLM @100 @100) + ] + , bgroup "200x200" + [ bench "hmatrix" $ nf (hmGemm hm200) hm200 + , bench "linear-massiv" $ nf (matMulP (mkMatLM @200 @200)) (mkMatLM @200 @200) + , bench "lm-parallel" $ nf (matMulPPar (mkMatLM @200 @200)) (mkMatLM @200 @200) + ] + , bgroup "500x500" + [ bench "hmatrix" $ nf (hmGemm hm500) hm500 + , bench "linear-massiv" $ nf (matMulP (mkMatLM @500 @500)) (mkMatLM @500 @500) + , bench "lm-parallel" $ nf (matMulPPar (mkMatLM @500 @500)) (mkMatLM @500 @500) + ] + ] + , bgroup "dot" + [ bgroup "4" + [ bench "linear" $ nf (LMet.dot linV4a) linV4b + , bench "hmatrix" $ nf (hmDot hv4) hv4 + , bench "linear-massiv" $ nf (dotP (mkVecLM @4)) (mkVecLM @4) + ] + , bgroup "100" + [ bench "hmatrix" $ nf (hmDot hv100) hv100 + , bench "linear-massiv" $ nf (dotP (mkVecLM @100)) (mkVecLM @100) + ] + , bgroup "1000" + [ bench "hmatrix" $ nf (hmDot hv1000) hv1000 + , bench "linear-massiv" $ nf (dotP (mkVecLM @1000)) (mkVecLM @1000) + ] + ] + , bgroup "matvec" + [ bgroup "4" + [ bench "linear" $ nf (linM44a LM.!*) linV4a + , bench "hmatrix" $ nf (hmMatvec hm4) hv4 + , bench "linear-massiv" $ nf (matvecP (mkMatLM @4 @4)) (mkVecLM @4) + ] + , bgroup "50" + [ bench "hmatrix" $ nf (hmMatvec hm50) hv50 + , bench "linear-massiv" $ nf (matvecP (mkMatLM @50 @50)) (mkVecLM @50) + ] + , bgroup "100" + [ bench "hmatrix" $ nf (hmMatvec hm100) hv100 + , bench "linear-massiv" $ nf (matvecP (mkMatLM @100 @100)) (mkVecLM @100) + ] + ] + , bgroup "luSolve" + [ bgroup "10x10" + [ bench "hmatrix" $ nf (hmLinearSolve dd10) hv10 + , bench "linear-massiv" $ nf (luSolveP (mkDDLM @10)) (mkVecLM @10) + , bench "lm-generic" $ nf (luSolve (mkDDLM @10)) (mkVecLM @10) + ] + , bgroup "50x50" + [ bench "hmatrix" $ nf (hmLinearSolve dd50) hv50 + , bench "linear-massiv" $ nf (luSolveP (mkDDLM @50)) (mkVecLM @50) + , bench "lm-generic" $ nf (luSolve (mkDDLM @50)) (mkVecLM @50) + ] + , bgroup "100x100" + [ bench "hmatrix" $ nf (hmLinearSolve dd100) hv100 + , bench "linear-massiv" $ nf (luSolveP (mkDDLM @100)) (mkVecLM @100) + , bench "lm-generic" $ nf (luSolve (mkDDLM @100)) (mkVecLM @100) + ] + ] + , bgroup "choleskySolve" + [ bgroup "10x10" + [ bench "hmatrix" $ nf (hmCholSolve spd10) hv10 + , bench "linear-massiv" $ nf (choleskySolveP (mkSPDLM @10)) (mkVecLM @10) + , bench "lm-generic" $ nf (choleskySolve (mkSPDLM @10)) (mkVecLM @10) + ] + , bgroup "50x50" + [ bench "hmatrix" $ nf (hmCholSolve spd50) hv50 + , bench "linear-massiv" $ nf (choleskySolveP (mkSPDLM @50)) (mkVecLM @50) + , bench "lm-generic" $ nf (choleskySolve (mkSPDLM @50)) (mkVecLM @50) + ] + , bgroup "100x100" + [ bench "hmatrix" $ nf (hmCholSolve spd100) hv100 + , bench "linear-massiv" $ nf (choleskySolveP (mkSPDLM @100)) (mkVecLM @100) + , bench "lm-generic" $ nf (choleskySolve (mkSPDLM @100)) (mkVecLM @100) + ] + ] + , bgroup "QR" + [ bgroup "10x10" + [ bench "hmatrix" $ nf H.qr hm10 + , bench "linear-massiv" $ nf qrP (mkMatLM @10 @10) + , bench "lm-generic" $ nf qr (mkMatLM @10 @10) + ] + , bgroup "50x50" + [ bench "hmatrix" $ nf H.qr hm50 + , bench "linear-massiv" $ nf qrP (mkMatLM @50 @50) + , bench "lm-generic" $ nf qr (mkMatLM @50 @50) + ] + , bgroup "100x100" + [ bench "hmatrix" $ nf H.qr hm100 + , bench "linear-massiv" $ nf qrP (mkMatLM @100 @100) + , bench "lm-generic" $ nf qr (mkMatLM @100 @100) + ] + ] + , bgroup "eigenSH" + [ bgroup "10x10" + [ bench "hmatrix" $ nf (H.eigSH . H.trustSym) spd10 + , bench "linear-massiv" $ nf (\a -> symmetricEigenP a 200 1e-12) (mkSPDLM @10) + , bench "lm-generic" $ nf (\a -> symmetricEigen a 200 1e-12) (mkSPDLM @10) + ] + , bgroup "50x50" + [ bench "hmatrix" $ nf (H.eigSH . H.trustSym) spd50 + , bench "linear-massiv" $ nf (\a -> symmetricEigenP a 500 1e-12) (mkSPDLM @50) + , bench "lm-generic" $ nf (\a -> symmetricEigen a 500 1e-12) (mkSPDLM @50) + ] + , bgroup "100x100" + [ bench "hmatrix" $ nf (H.eigSH . H.trustSym) spd100 + , bench "linear-massiv" $ nf (\a -> symmetricEigenP a 1000 1e-12) (mkSPDLM @100) + , bench "lm-parallel" $ nf (\a -> symmetricEigenPPar a 1000 1e-12) (mkSPDLM @100) + ] + , bgroup "200x200" + [ bench "hmatrix" $ nf (H.eigSH . H.trustSym) spd200 + , bench "linear-massiv" $ nf (\a -> symmetricEigenP a 2000 1e-12) (mkSPDLM @200) + ] + , bgroup "500x500" + [ bench "hmatrix" $ nf (H.eigSH . H.trustSym) spd500 + , bench "linear-massiv" $ nf (\a -> symmetricEigenP a 5000 1e-12) (mkSPDLM @500) + ] + ] + , bgroup "eigenSH-breakdown" + [ bgroup "200x200" + [ bench "tridiagP-only" $ nf tridiagonalizeP (mkSPDLM @200) + , bench "full-eigenP" $ nf (\a -> symmetricEigenP a 2000 1e-12) (mkSPDLM @200) + ] + , bgroup "500x500" + [ bench "tridiagP-only" $ nf tridiagonalizeP (mkSPDLM @500) + , bench "full-eigenP" $ nf (\a -> symmetricEigenP a 5000 1e-12) (mkSPDLM @500) + ] + ] + , bgroup "eigenSH-DC" + [ bgroup "50x50" + [ bench "QR" $ nf (\a -> symmetricEigenP a 500 1e-12) (mkSPDLM @50) + , bench "D&C" $ nf (\a -> symmetricEigenPDC a 1e-12) (mkSPDLM @50) + ] + , bgroup "100x100" + [ bench "QR" $ nf (\a -> symmetricEigenP a 1000 1e-12) (mkSPDLM @100) + , bench "D&C" $ nf (\a -> symmetricEigenPDC a 1e-12) (mkSPDLM @100) + ] + , bgroup "200x200" + [ bench "QR" $ nf (\a -> symmetricEigenP a 2000 1e-12) (mkSPDLM @200) + , bench "D&C" $ nf (\a -> symmetricEigenPDC a 1e-12) (mkSPDLM @200) + ] + , bgroup "500x500" + [ bench "QR" $ nf (\a -> symmetricEigenP a 5000 1e-12) (mkSPDLM @500) + , bench "D&C" $ nf (\a -> symmetricEigenPDC a 1e-12) (mkSPDLM @500) + ] + ] + , bgroup "SVD" + [ bgroup "10x10" + [ bench "hmatrix" $ nf H.svd hm10 + , bench "linear-massiv" $ nf svdP (mkMatLM @10 @10) + , bench "lm-generic" $ nf svd (mkMatLM @10 @10) + ] + , bgroup "50x50" + [ bench "hmatrix" $ nf H.svd hm50 + , bench "linear-massiv" $ nf svdP (mkMatLM @50 @50) + , bench "lm-generic" $ nf svd (mkMatLM @50 @50) + ] + , bgroup "100x100" + [ bench "hmatrix" $ nf H.svd hm100 + , bench "linear-massiv" $ nf svdP (mkMatLM @100 @100) + , bench "lm-gk" $ nf svdGKP (mkMatLM @100 @100) + ] + , bgroup "200x200" + [ bench "hmatrix" $ nf H.svd hm200 + , bench "linear-massiv" $ nf svdP (mkMatLM @200 @200) + , bench "lm-gk" $ nf svdGKP (mkMatLM @200 @200) + ] + , bgroup "500x500" + [ bench "hmatrix" $ nf H.svd hm500 + , bench "linear-massiv" $ nf svdP (mkMatLM @500 @500) + , bench "lm-gk" $ nf svdGKP (mkMatLM @500 @500) + ] + ] + ] diff --git a/bench/Bench/BLAS.hs b/bench/Bench/BLAS.hs new file mode 100644 --- /dev/null +++ b/bench/Bench/BLAS.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Bench.BLAS (blasBenchmarks) where + +import Criterion.Main +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Comp(..)) +import Control.DeepSeq () +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level1 (dotP) +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvecP) +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMulP, matMulComp) + +-- linear library imports for comparison +import Linear.V4 (V4(..)) +import Linear.Matrix ((!*!), (!*)) +import qualified Linear.Metric as LM +import Linear.V4 () + +-- Massiv matrix generators +mkMatP :: forall m n. (KnownNat m, KnownNat n) => Matrix m n M.P Double +mkMatP = makeMatrix @m @n @M.P $ \i j -> fromIntegral (i * 7 + j * 3 + 1) / 100.0 + +mkVecP :: forall n. KnownNat n => Vector n M.P Double +mkVecP = makeVector @n @M.P $ \i -> fromIntegral (i + 1) / 10.0 + +-- linear library 4x4 matrices for comparison +linearM44 :: V4 (V4 Double) +linearM44 = V4 (V4 1 2 3 4) + (V4 5 6 7 8) + (V4 9 10 11 12) + (V4 13 14 15 16) + +linearM44b :: V4 (V4 Double) +linearM44b = V4 (V4 17 18 19 20) + (V4 21 22 23 24) + (V4 25 26 27 28) + (V4 29 30 31 32) + +linearV4 :: V4 Double +linearV4 = V4 1 2 3 4 + +linearV4b :: V4 Double +linearV4b = V4 5 6 7 8 + +blasBenchmarks :: [Benchmark] +blasBenchmarks = + [ bgroup "gemm" + [ -- Small sizes: compare linear vs massiv + bgroup "4x4" + [ bench "linear-V4" $ nf (linearM44 !*!) linearM44b + , bench "massiv-P" $ nf (uncurry (matMulP @4 @4 @4)) (mkMatP @4 @4, mkMatP @4 @4) + ] + , bgroup "10x10" + [ bench "P/Seq" $ nf (uncurry (matMulP @10 @10 @10)) (mkMatP @10 @10, mkMatP @10 @10) + ] + , bgroup "50x50" + [ bench "P/Seq" $ nf (uncurry (matMulP @50 @50 @50)) (mkMatP @50 @50, mkMatP @50 @50) + , bench "P/Par" $ nf (uncurry (matMulComp @50 @50 @50 Par)) (mkMatP @50 @50, mkMatP @50 @50) + ] + , bgroup "100x100" + [ bench "P/Seq" $ nf (uncurry (matMulP @100 @100 @100)) (mkMatP @100 @100, mkMatP @100 @100) + , bench "P/Par" $ nf (uncurry (matMulComp @100 @100 @100 Par)) (mkMatP @100 @100, mkMatP @100 @100) + ] + , bgroup "200x200" + [ bench "P/Seq" $ nf (uncurry (matMulP @200 @200 @200)) (mkMatP @200 @200, mkMatP @200 @200) + , bench "P/Par" $ nf (uncurry (matMulComp @200 @200 @200 Par)) (mkMatP @200 @200, mkMatP @200 @200) + ] + -- Representation comparison at 100x100 + , bgroup "repr-100x100" + [ bench "P" $ nf (uncurry (matMulP @100 @100 @100)) (mkMatP @100 @100, mkMatP @100 @100) + ] + ] + , bgroup "dot" + [ bgroup "4" + [ bench "linear-V4" $ nf (LM.dot linearV4) linearV4b + , bench "massiv-P" $ nf (uncurry (dotP @4)) (mkVecP @4, mkVecP @4) + ] + , bgroup "100" + [ bench "P" $ nf (uncurry (dotP @100)) (mkVecP @100, mkVecP @100) + ] + , bgroup "1000" + [ bench "P" $ nf (uncurry (dotP @1000)) (mkVecP @1000, mkVecP @1000) + ] + , bgroup "10000" + [ bench "P" $ nf (uncurry (dotP @10000)) (mkVecP @10000, mkVecP @10000) + ] + ] + , bgroup "matvec" + [ bgroup "4" + [ bench "linear-V4" $ nf (linearM44 !*) linearV4 + , bench "massiv-P" $ nf (uncurry (matvecP @4 @4)) (mkMatP @4 @4, mkVecP @4) + ] + , bgroup "50" + [ bench "P" $ nf (uncurry (matvecP @50 @50)) (mkMatP @50 @50, mkVecP @50) + ] + , bgroup "100" + [ bench "P" $ nf (uncurry (matvecP @100 @100)) (mkMatP @100 @100, mkVecP @100) + ] + ] + ] diff --git a/bench/Bench/Eigen.hs b/bench/Bench/Eigen.hs new file mode 100644 --- /dev/null +++ b/bench/Bench/Eigen.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Bench.Eigen (eigenBenchmarks) where + +import Criterion.Main +import qualified Data.Massiv.Array as M +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Eigen.Symmetric (symmetricEigen, jacobiEigen) +import Numeric.LinearAlgebra.Massiv.Eigen.SVD (svd) + +mkSPDP :: forall n. KnownNat n => Matrix n n M.P Double +mkSPDP = + let nn = dimVal @n + b = makeMatrix @n @n @M.P $ \i j -> fromIntegral (i * nn + j + 1) / fromIntegral (nn * nn) + in makeMatrix @n @n @M.P $ \i j -> + foldl' (\acc k -> acc + (b ! (i, k)) * (b ! (j, k))) (if i == j then 1 else 0) [0..nn-1] + +mkMatP :: forall m n. (KnownNat m, KnownNat n) => Matrix m n M.P Double +mkMatP = makeMatrix @m @n @M.P $ \i j -> fromIntegral (i * 7 + j * 3 + 1) / 100.0 + +eigenBenchmarks :: [Benchmark] +eigenBenchmarks = + [ bgroup "symmetricEigen" + [ bench "10x10/P" $ nf (\a -> symmetricEigen a 200 1e-12) (mkSPDP @10) + , bench "20x20/P" $ nf (\a -> symmetricEigen a 500 1e-12) (mkSPDP @20) + , bench "50x50/P" $ nf (\a -> symmetricEigen a 500 1e-12) (mkSPDP @50) + ] + , bgroup "jacobiEigen" + [ bench "10x10/P" $ nf (\a -> jacobiEigen a 50 1e-12) (mkSPDP @10) + , bench "20x20/P" $ nf (\a -> jacobiEigen a 50 1e-12) (mkSPDP @20) + ] + , bgroup "svd" + [ bench "10x10/P" $ nf svd (mkMatP @10 @10) + , bench "20x20/P" $ nf svd (mkMatP @20 @20) + , bench "50x50/P" $ nf svd (mkMatP @50 @50) + ] + ] diff --git a/bench/Bench/Orthogonal.hs b/bench/Bench/Orthogonal.hs new file mode 100644 --- /dev/null +++ b/bench/Bench/Orthogonal.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Bench.Orthogonal (orthogonalBenchmarks) where + +import Criterion.Main +import qualified Data.Massiv.Array as M +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Orthogonal.QR (qr, qrGivens) + +mkMatP :: forall m n. (KnownNat m, KnownNat n) => Matrix m n M.P Double +mkMatP = makeMatrix @m @n @M.P $ \i j -> fromIntegral (i * 7 + j * 3 + 1) / 100.0 + +orthogonalBenchmarks :: [Benchmark] +orthogonalBenchmarks = + [ bgroup "qr-householder" + [ bench "10x10/P" $ nf qr (mkMatP @10 @10) + , bench "50x50/P" $ nf qr (mkMatP @50 @50) + , bench "100x100/P" $ nf qr (mkMatP @100 @100) + ] + , bgroup "qr-givens" + [ bench "10x10/P" $ nf qrGivens (mkMatP @10 @10) + , bench "50x50/P" $ nf qrGivens (mkMatP @50 @50) + ] + ] diff --git a/bench/Bench/Parallel.hs b/bench/Bench/Parallel.hs new file mode 100644 --- /dev/null +++ b/bench/Bench/Parallel.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | Benchmarks specifically for measuring parallelism scalability. +-- +-- Runs matrix multiplication at various sizes with explicit thread counts +-- using massiv's ParN constructor. +module Bench.Parallel (parallelBenchmarks) where + +import Criterion.Main +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Comp(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMulComp) + +mkMatP :: forall m n. (KnownNat m, KnownNat n) => Matrix m n M.P Double +mkMatP = makeMatrix @m @n @M.P $ \i j -> fromIntegral (i * 7 + j * 3 + 1) / 100.0 + +-- | Benchmarks measuring how performance scales with thread count. +parallelBenchmarks :: [Benchmark] +parallelBenchmarks = + [ bgroup "scaling-gemm" + [ bgroup "100x100" + [ bench "Seq" $ nf (uncurry (matMulComp @100 @100 @100 Seq)) (mkMatP @100 @100, mkMatP @100 @100) + , bench "Par" $ nf (uncurry (matMulComp @100 @100 @100 Par)) (mkMatP @100 @100, mkMatP @100 @100) + , bench "ParN-1" $ nf (uncurry (matMulComp @100 @100 @100 (ParN 1))) (mkMatP @100 @100, mkMatP @100 @100) + , bench "ParN-2" $ nf (uncurry (matMulComp @100 @100 @100 (ParN 2))) (mkMatP @100 @100, mkMatP @100 @100) + , bench "ParN-4" $ nf (uncurry (matMulComp @100 @100 @100 (ParN 4))) (mkMatP @100 @100, mkMatP @100 @100) + , bench "ParN-8" $ nf (uncurry (matMulComp @100 @100 @100 (ParN 8))) (mkMatP @100 @100, mkMatP @100 @100) + , bench "ParN-16" $ nf (uncurry (matMulComp @100 @100 @100 (ParN 16))) (mkMatP @100 @100, mkMatP @100 @100) + , bench "ParN-20" $ nf (uncurry (matMulComp @100 @100 @100 (ParN 20))) (mkMatP @100 @100, mkMatP @100 @100) + ] + , bgroup "200x200" + [ bench "Seq" $ nf (uncurry (matMulComp @200 @200 @200 Seq)) (mkMatP @200 @200, mkMatP @200 @200) + , bench "Par" $ nf (uncurry (matMulComp @200 @200 @200 Par)) (mkMatP @200 @200, mkMatP @200 @200) + , bench "ParN-1" $ nf (uncurry (matMulComp @200 @200 @200 (ParN 1))) (mkMatP @200 @200, mkMatP @200 @200) + , bench "ParN-2" $ nf (uncurry (matMulComp @200 @200 @200 (ParN 2))) (mkMatP @200 @200, mkMatP @200 @200) + , bench "ParN-4" $ nf (uncurry (matMulComp @200 @200 @200 (ParN 4))) (mkMatP @200 @200, mkMatP @200 @200) + , bench "ParN-8" $ nf (uncurry (matMulComp @200 @200 @200 (ParN 8))) (mkMatP @200 @200, mkMatP @200 @200) + , bench "ParN-16" $ nf (uncurry (matMulComp @200 @200 @200 (ParN 16))) (mkMatP @200 @200, mkMatP @200 @200) + , bench "ParN-20" $ nf (uncurry (matMulComp @200 @200 @200 (ParN 20))) (mkMatP @200 @200, mkMatP @200 @200) + ] + ] + ] diff --git a/bench/Bench/Solve.hs b/bench/Bench/Solve.hs new file mode 100644 --- /dev/null +++ b/bench/Bench/Solve.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Bench.Solve (solveBenchmarks) where + +import Criterion.Main +import qualified Data.Massiv.Array as M +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Solve.LU (lu, luSolve) +import Numeric.LinearAlgebra.Massiv.Solve.Cholesky (cholesky, choleskySolve) + +-- Diagonally dominant matrix for LU +mkMatP :: forall n. KnownNat n => Matrix n n M.P Double +mkMatP = makeMatrix @n @n @M.P $ \i j -> + fromIntegral (i * 7 + j * 3 + 1) / 100.0 + if i == j then fromIntegral (dimVal @n) else 0 + +-- SPD matrix: A = B^T B + nI +mkSPDP :: forall n. KnownNat n => Matrix n n M.P Double +mkSPDP = + let nn = dimVal @n + b = makeMatrix @n @n @M.P $ \i j -> fromIntegral (i * nn + j + 1) / fromIntegral (nn * nn) + in makeMatrix @n @n @M.P $ \i j -> + foldl' (\acc k -> acc + (b ! (i, k)) * (b ! (j, k))) (if i == j then 1 else 0) [0..nn-1] + +mkVecP :: forall n. KnownNat n => Vector n M.P Double +mkVecP = makeVector @n @M.P $ \i -> fromIntegral (i + 1) + +solveBenchmarks :: [Benchmark] +solveBenchmarks = + [ bgroup "lu" + [ bench "10x10/P" $ nf lu (mkMatP @10) + , bench "50x50/P" $ nf lu (mkMatP @50) + , bench "100x100/P" $ nf lu (mkMatP @100) + ] + , bgroup "luSolve" + [ bench "10x10/P" $ nf (uncurry luSolve) (mkMatP @10, mkVecP @10) + , bench "50x50/P" $ nf (uncurry luSolve) (mkMatP @50, mkVecP @50) + , bench "100x100/P" $ nf (uncurry luSolve) (mkMatP @100, mkVecP @100) + ] + , bgroup "cholesky" + [ bench "10x10/P" $ nf cholesky (mkSPDP @10) + , bench "50x50/P" $ nf cholesky (mkSPDP @50) + , bench "100x100/P" $ nf cholesky (mkSPDP @100) + ] + , bgroup "choleskySolve" + [ bench "10x10/P" $ nf (uncurry choleskySolve) (mkSPDP @10, mkVecP @10) + , bench "50x50/P" $ nf (uncurry choleskySolve) (mkSPDP @50, mkVecP @50) + , bench "100x100/P" $ nf (uncurry choleskySolve) (mkSPDP @100, mkVecP @100) + ] + ] diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,18 @@ +module Main (main) where + +import Criterion.Main + +import Bench.BLAS (blasBenchmarks) +import Bench.Solve (solveBenchmarks) +import Bench.Orthogonal (orthogonalBenchmarks) +import Bench.Eigen (eigenBenchmarks) +import Bench.Parallel (parallelBenchmarks) + +main :: IO () +main = defaultMain + [ bgroup "BLAS" blasBenchmarks + , bgroup "Solve" solveBenchmarks + , bgroup "Orthogonal" orthogonalBenchmarks + , bgroup "Eigen" eigenBenchmarks + , bgroup "Parallel" parallelBenchmarks + ] diff --git a/linear-massiv.cabal b/linear-massiv.cabal new file mode 100644 --- /dev/null +++ b/linear-massiv.cabal @@ -0,0 +1,159 @@ +cabal-version: 3.0 +name: linear-massiv +version: 0.1.0.0 +synopsis: Type-safe numerical linear algebra backed by massiv arrays +description: + Native Haskell implementations of algorithms from Golub & Van Loan's + "Matrix Computations" (4th ed.) using massiv arrays as the backing store, + with compile-time dimensional conformance via GHC type-level naturals. + Extends the linear library's typeclasses for integration with existing code. + . + Co-authored by Claude Opus (Anthropic). This code should be considered a + derived work of the various algorithmic examples and reference + implementations drawn upon during development, including but not limited + to LAPACK, OpenBLAS, and GVL4. +license: BSD-3-Clause +build-type: Simple + +library + hs-source-dirs: src + exposed-modules: + Numeric.LinearAlgebra.Massiv + Numeric.LinearAlgebra.Massiv.Types + Numeric.LinearAlgebra.Massiv.Internal + Numeric.LinearAlgebra.Massiv.Norms + Numeric.LinearAlgebra.Massiv.BLAS.Level1 + Numeric.LinearAlgebra.Massiv.BLAS.Level2 + Numeric.LinearAlgebra.Massiv.BLAS.Level3 + Numeric.LinearAlgebra.Massiv.Solve.Triangular + Numeric.LinearAlgebra.Massiv.Solve.LU + Numeric.LinearAlgebra.Massiv.Solve.Cholesky + Numeric.LinearAlgebra.Massiv.Solve.Banded + Numeric.LinearAlgebra.Massiv.Orthogonal.Householder + Numeric.LinearAlgebra.Massiv.Orthogonal.Givens + Numeric.LinearAlgebra.Massiv.Orthogonal.QR + Numeric.LinearAlgebra.Massiv.Orthogonal.LeastSquares + Numeric.LinearAlgebra.Massiv.Eigen.Power + Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg + Numeric.LinearAlgebra.Massiv.Eigen.Schur + Numeric.LinearAlgebra.Massiv.Eigen.Symmetric + Numeric.LinearAlgebra.Massiv.Eigen.SVD + Numeric.LinearAlgebra.Massiv.Internal.Kernel + Numeric.LinearAlgebra.Massiv.Linear + build-depends: + base >= 4.16 && < 5, + ghc-prim, + massiv >= 1.0 && < 2, + linear >= 1.21 && < 2, + vector >= 0.12 && < 1, + primitive >= 0.7 && < 1, + deepseq >= 1.4 && < 2 + default-language: GHC2021 + default-extensions: + DataKinds + TypeFamilies + ScopedTypeVariables + TypeApplications + TypeOperators + GADTs + RankNTypes + FlexibleContexts + FlexibleInstances + MultiParamTypeClasses + StandaloneDeriving + DerivingStrategies + ghc-options: -Wall -O2 + +test-suite linear-massiv-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + hs-source-dirs: test + other-modules: + Test.Types + Test.Residuals + Test.BLAS + Test.Solve + Test.Orthogonal + Test.Eigen + Test.Norms + build-depends: + base, + linear-massiv, + massiv, + tasty >= 1.4, + tasty-hunit >= 0.10, + tasty-quickcheck >= 0.10, + QuickCheck >= 2.14, + linear, + vector + default-language: GHC2021 + default-extensions: + DataKinds + TypeFamilies + ScopedTypeVariables + TypeApplications + TypeOperators + GADTs + RankNTypes + FlexibleContexts + FlexibleInstances + MultiParamTypeClasses + ghc-options: -Wall -threaded -rtsopts + +benchmark linear-massiv-bench + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench + other-modules: + Bench.BLAS + Bench.Solve + Bench.Orthogonal + Bench.Eigen + Bench.Parallel + build-depends: + base, + linear-massiv, + massiv, + criterion >= 1.5, + deepseq, + linear + default-language: GHC2021 + default-extensions: + DataKinds + TypeFamilies + ScopedTypeVariables + TypeApplications + TypeOperators + GADTs + RankNTypes + FlexibleContexts + FlexibleInstances + MultiParamTypeClasses + ghc-options: -Wall -O2 -threaded -rtsopts "-with-rtsopts=-N" + +benchmark linear-massiv-comparison + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench-comparison + build-depends: + base, + linear-massiv, + massiv, + hmatrix >= 0.20, + linear, + criterion >= 1.5, + deepseq, + vector + default-language: GHC2021 + default-extensions: + DataKinds + TypeFamilies + ScopedTypeVariables + TypeApplications + TypeOperators + GADTs + RankNTypes + FlexibleContexts + FlexibleInstances + MultiParamTypeClasses + ghc-options: -Wall -O2 -threaded -rtsopts "-with-rtsopts=-N" diff --git a/src/Numeric/LinearAlgebra/Massiv.hs b/src/Numeric/LinearAlgebra/Massiv.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv.hs @@ -0,0 +1,236 @@ +-- | +-- Module : Numeric.LinearAlgebra.Massiv +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- @linear-massiv@: Type-safe numerical linear algebra backed by +-- arrays. +-- +-- This library provides native Haskell implementations of algorithms from: +-- +-- * Golub, G. H., & Van Loan, C. F. (2013). /Matrix Computations/ (4th ed.). +-- Johns Hopkins University Press. ISBN 978-1-4214-0794-4. +-- +-- referred to throughout as __GVL4__. +-- +-- = Derived Work Attribution +-- +-- This code was co-authored by Claude Opus (Anthropic) and should be +-- considered a derived work of the various algorithmic examples and +-- reference implementations drawn upon during development, including but +-- not limited to: +-- +-- * __LAPACK__ (Linear Algebra PACKage) — Anderson, E. et al. (1999). +-- /LAPACK Users' Guide/, 3rd ed., SIAM. The LAPACK testing methodology, +-- algorithm structures, and numerical stability techniques informed much +-- of the implementation. +-- +-- * __OpenBLAS__ — Xianyi, Z., Qian, W., and Yunquan, Z. (2011--). +-- The tiled GEMM micro-kernel architecture, cache-blocking strategies, +-- and SIMD vectorisation patterns were inspired by OpenBLAS. +-- +-- * __GVL4__ — The primary algorithmic reference, as noted above. +-- +-- * __Higham__ — Higham, N. J. (2002). /Accuracy and Stability of Numerical +-- Algorithms/, 2nd ed., SIAM. Error analysis and numerical stability +-- frameworks. +-- +-- = Design Principles +-- +-- 1. __Type-level dimensional safety__: Matrix dimensions are tracked in the +-- type system via GHC @DataKinds@ and @KnownNat@ constraints. Dimensionally +-- incorrect operations (e.g., multiplying an \(m \times k\) matrix by an +-- \(n \times p\) matrix where \(k \neq n\)) are rejected at compile time. +-- +-- 2. __Representation polymorphism__: All operations are parametric over +-- massiv's array representation @r@ (e.g., @P@ for primitive, @U@ for +-- unboxed, @S@ for storable, @B@ for boxed), constrained by +-- @'Data.Massiv.Array.Manifest' r e@. Users choose the representation at +-- the call site. +-- +-- 3. __Parallelism via massiv__: Operations that construct arrays via +-- @makeArray@ inherit massiv's computation strategies. Use 'matMulComp' +-- with @Par@ or @ParN n@ for parallel matrix multiplication. +-- +-- 4. __No FFI__: All algorithms are pure Haskell, enabling portability and +-- auditability. Benchmarks compare performance across massiv representations +-- and parallelism strategies. +-- +-- = Internal Architecture: Two-Layer Design +-- +-- @linear-massiv@ uses a two-layer architecture that separates the type-safe +-- public API from the performance-critical internal representation: +-- +-- * __Public layer__: 'Matrix' and 'Vector' are @newtype@ wrappers around +-- massiv's @Array r Ix2 e@, providing compile-time dimension checking via +-- phantom @Nat@ parameters and representation polymorphism via @r@. +-- +-- * __Internal layer__: Performance-critical inner loops (GEMM, QR, +-- tridiagonalisation, SVD, etc.) unwrap the massiv array to a raw +-- @ByteArray#@ \/ @MutableByteArray#@ and operate directly via GHC primops, +-- including @DoubleX4#@ AVX2 SIMD instructions compiled through the LLVM 17 +-- backend. Functions receive @(ByteArray, offset, stride)@ triples, enabling +-- zero-copy submatrix views for panel factorisations. +-- +-- This separation is essential for performance. Benchmarks (Round 3 of the +-- accompanying report) showed that massiv's per-element @M.readM@\/@M.write_@ +-- abstraction layer imposed a 240–330× penalty on BLAS operations relative to +-- direct primop access, even though the underlying memory layout is identical. +-- The raw primop layer eliminates this overhead while the @newtype@ wrapper +-- preserves type safety at the API boundary. +-- +-- == Why not @vector-sized@ or @linear@'s @V@? +-- +-- The @@ package +-- provides an @Unbox (Vector n a)@ instance that stores +-- @Vector m (Vector n Double)@ as a contiguous flat @ByteArray@ of @m × n@ +-- doubles. While the __memory layout__ is correct, contiguous memory alone is +-- insufficient for high-performance numerical kernels: +-- +-- * __Per-element typeclass dispatch__: Every access goes through +-- @basicUnsafeRead@ \/ @basicUnsafeWrite@ of the @Unbox@ data family. +-- Reading element @(i, j)@ requires indexing the outer vector to obtain a +-- @Vector n Double@ (constructing an intermediate slice), then indexing that. +-- @linear-massiv@ computes @off + i * stride + j@ and issues a single +-- @readDoubleArray#@ primop. +-- +-- * __No SIMD access__: The 4×8 GEMM micro-kernel loads 4 consecutive doubles +-- via @indexDoubleArrayAsDoubleX4# ba# (off# +# i#)@—a direct 256-bit AVX2 +-- load from a computed byte offset. The @Unbox@ typeclass does not expose +-- the underlying @ByteArray#@, and GHC cannot optimise through the data +-- family indirection to produce equivalent code. +-- +-- * __No mutable primop access__: In-place factorisations (LU, QR, +-- tridiagonalisation, bidiagonalisation) require @writeDoubleArray#@ on +-- @MutableByteArray#@ with computed offsets. The @MVector@ abstraction +-- interposes allocation and function calls that prevent the tight unboxed +-- loops needed for competitive performance. +-- +-- * __No zero-copy submatrix views__: Panel factorisations pass +-- @(ByteArray, offset, stride)@ triples to kernels, enabling zero-copy views +-- into submatrices. @Vector n a@ does not naturally express "this row starts +-- at byte offset X in a larger backing array." +-- +-- The @@ library's @V n a@ +-- uses @Vector@ from the @vector@ package internally and is designed for small +-- fixed-size vectors (V2–V4) where GHC can fully unbox everything. At +-- @n = 100–500@, @V n (V m Double)@ would be a vector-of-vectors with per-row +-- indirection—catastrophic for cache locality and SIMD vectorisation. +-- @linear-massiv@ provides conversion functions ('fromLinearV', 'fromV2', etc.) +-- in "Numeric.LinearAlgebra.Massiv.Linear" for interoperability. +-- +-- = Quick Start +-- +-- @ +-- import Numeric.LinearAlgebra.Massiv +-- import qualified Data.Massiv.Array as M +-- +-- -- Create a 3x3 matrix (Primitive representation, Double elements) +-- let a = 'makeMatrix' \@3 \@3 \@M.P $ \\i j -> +-- fromIntegral (i * 3 + j + 1) :: Double +-- +-- -- QR factorization +-- let (q, r) = 'qr' a +-- +-- -- Solve Ax = b via LU +-- let b = 'makeVector' \@3 \@M.P $ \\i -> fromIntegral (i + 1) :: Double +-- let x = 'luSolve' a b +-- @ +-- +-- = Module Organisation +-- +-- == Core types and construction +-- +-- * "Numeric.LinearAlgebra.Massiv.Types" — 'Matrix', 'Vector' newtypes with +-- phantom dimension parameters +-- * "Numeric.LinearAlgebra.Massiv.Internal" — Unsafe constructors, dimension +-- reification, array creation helpers +-- +-- == BLAS-like operations (GVL4 Ch. 1) +-- +-- * "Numeric.LinearAlgebra.Massiv.BLAS.Level1" — Vector–vector: 'dot', 'axpy', 'scal', 'nrm2' +-- * "Numeric.LinearAlgebra.Massiv.BLAS.Level2" — Matrix–vector: 'gemv', 'matvec', 'ger' +-- * "Numeric.LinearAlgebra.Massiv.BLAS.Level3" — Matrix–matrix: 'gemm', 'matMul', 'transpose' +-- +-- == Direct solvers (GVL4 Chs. 3–4) +-- +-- * "Numeric.LinearAlgebra.Massiv.Solve.Triangular" — Forward\/back substitution +-- * "Numeric.LinearAlgebra.Massiv.Solve.LU" — LU with partial pivoting +-- * "Numeric.LinearAlgebra.Massiv.Solve.Cholesky" — Cholesky for SPD matrices +-- * "Numeric.LinearAlgebra.Massiv.Solve.Banded" — Band LU, band Cholesky, tridiagonal solver +-- +-- == Orthogonal factorizations (GVL4 Ch. 5) +-- +-- * "Numeric.LinearAlgebra.Massiv.Orthogonal.Householder" — Householder reflections +-- * "Numeric.LinearAlgebra.Massiv.Orthogonal.Givens" — Givens rotations +-- * "Numeric.LinearAlgebra.Massiv.Orthogonal.QR" — QR factorization (Householder and Givens) +-- * "Numeric.LinearAlgebra.Massiv.Orthogonal.LeastSquares" — Least squares via QR +-- +-- == Eigenvalue problems and SVD (GVL4 Chs. 7–8) +-- +-- * "Numeric.LinearAlgebra.Massiv.Eigen.Power" — Power, inverse, Rayleigh quotient iteration +-- * "Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg" — Hessenberg reduction +-- * "Numeric.LinearAlgebra.Massiv.Eigen.Schur" — Schur decomposition (QR algorithm) +-- * "Numeric.LinearAlgebra.Massiv.Eigen.Symmetric" — Symmetric eigenvalue (tridiagonal QR, Jacobi) +-- * "Numeric.LinearAlgebra.Massiv.Eigen.SVD" — Singular value decomposition +-- +-- == Norms and condition numbers (GVL4 Ch. 2) +-- +-- * "Numeric.LinearAlgebra.Massiv.Norms" — Frobenius, 1-, \(\infty\)-, and 2-norms +-- +-- == Integration with the @linear@ library +-- +-- * "Numeric.LinearAlgebra.Massiv.Linear" — Conversions to\/from @linear@'s @V@, @V2@, @V3@, @V4@ +module Numeric.LinearAlgebra.Massiv + ( -- * Core types + module Numeric.LinearAlgebra.Massiv.Types + -- * Construction helpers + , module Numeric.LinearAlgebra.Massiv.Internal + -- * BLAS operations + , module Numeric.LinearAlgebra.Massiv.BLAS.Level1 + , module Numeric.LinearAlgebra.Massiv.BLAS.Level2 + , module Numeric.LinearAlgebra.Massiv.BLAS.Level3 + -- * Direct solvers + , module Numeric.LinearAlgebra.Massiv.Solve.Triangular + , module Numeric.LinearAlgebra.Massiv.Solve.LU + , module Numeric.LinearAlgebra.Massiv.Solve.Cholesky + , module Numeric.LinearAlgebra.Massiv.Solve.Banded + -- * Orthogonal factorizations + , module Numeric.LinearAlgebra.Massiv.Orthogonal.Householder + , module Numeric.LinearAlgebra.Massiv.Orthogonal.Givens + , module Numeric.LinearAlgebra.Massiv.Orthogonal.QR + , module Numeric.LinearAlgebra.Massiv.Orthogonal.LeastSquares + -- * Eigenvalue problems + , module Numeric.LinearAlgebra.Massiv.Eigen.Power + , module Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg + , module Numeric.LinearAlgebra.Massiv.Eigen.Schur + , module Numeric.LinearAlgebra.Massiv.Eigen.Symmetric + , module Numeric.LinearAlgebra.Massiv.Eigen.SVD + -- * Norms + , module Numeric.LinearAlgebra.Massiv.Norms + -- * Linear integration + , module Numeric.LinearAlgebra.Massiv.Linear + ) where + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level1 +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 +import Numeric.LinearAlgebra.Massiv.Solve.Triangular +import Numeric.LinearAlgebra.Massiv.Solve.LU +import Numeric.LinearAlgebra.Massiv.Solve.Cholesky +import Numeric.LinearAlgebra.Massiv.Solve.Banded +import Numeric.LinearAlgebra.Massiv.Orthogonal.Householder +import Numeric.LinearAlgebra.Massiv.Orthogonal.Givens +import Numeric.LinearAlgebra.Massiv.Orthogonal.QR +import Numeric.LinearAlgebra.Massiv.Orthogonal.LeastSquares +import Numeric.LinearAlgebra.Massiv.Eigen.Power +import Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg +import Numeric.LinearAlgebra.Massiv.Eigen.Schur +import Numeric.LinearAlgebra.Massiv.Eigen.Symmetric +import Numeric.LinearAlgebra.Massiv.Eigen.SVD +import Numeric.LinearAlgebra.Massiv.Norms +import Numeric.LinearAlgebra.Massiv.Linear diff --git a/src/Numeric/LinearAlgebra/Massiv/BLAS/Level1.hs b/src/Numeric/LinearAlgebra/Massiv/BLAS/Level1.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/BLAS/Level1.hs @@ -0,0 +1,203 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.BLAS.Level1 +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = BLAS Level 1: Vector–Vector Operations +-- +-- This module provides type-safe, dimension-indexed wrappers around the +-- standard BLAS Level 1 kernels for vector–vector operations. Every +-- function carries the vector length @n@ as a phantom type-level natural, +-- so dimension mismatches are caught at compile time rather than at run +-- time. +-- +-- The algorithms implemented here correspond to the elementary building +-- blocks described in: +-- +-- * Golub, G. H. & Van Loan, C. F. (2013). /Matrix Computations/, +-- 4th edition (GVL4). Johns Hopkins University Press. +-- __Chapter 1, Section 1.1__, pp. 4–8. +-- +-- Specifically: +-- +-- * __Algorithm 1.1.1__ (p. 4) — Inner product (dot product). +-- Given vectors \(x, y \in \mathbb{R}^{n}\), compute +-- \(c = x^{T} y = \sum_{i=1}^{n} x_i y_i\). +-- +-- * __Algorithm 1.1.2 (Saxpy)__ (p. 4) — Scalar \(\alpha\) times +-- vector \(x\) plus vector \(y\): +-- \(y \leftarrow \alpha x + y\). +-- This is the fundamental vector-update operation upon which the +-- higher-level BLAS Level 2 and Level 3 routines are built. +-- +-- Additionally the module exposes the common vector norms +-- (\(\lVert \cdot \rVert_1\) and \(\lVert \cdot \rVert_2\)) and +-- scalar–vector scaling, which together form the complete Level 1 +-- BLAS surface. +-- +-- == Complexity +-- +-- All operations in this module are \(O(n)\) in the vector length. +module Numeric.LinearAlgebra.Massiv.BLAS.Level1 + ( -- * Dot product (Algorithm 1.1.1, GVL4 p. 4) + dot + , dotP + -- * Scalar–vector operations (Algorithm 1.1.2, GVL4 p. 4) + , scal + , axpy + -- * Vector norms (GVL4 Section 1.1, pp. 4–8) + , nrm2 + , asum + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (unwrapByteArray, unwrapByteArrayOffset) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Internal.Kernel (rawDot) + +-- | Inner (dot) product of two vectors. +-- +-- __GVL4 Reference:__ Algorithm 1.1.1, p. 4. +-- +-- Given \(x, y \in \mathbb{R}^{n}\), computes the scalar +-- +-- \[ +-- c \;=\; x^{T} y \;=\; \sum_{i=1}^{n} x_i \, y_i +-- \] +-- +-- ==== Type-safety guarantees +-- +-- Both vectors carry the same compile-time dimension @n@, so a +-- length mismatch is a type error. +-- +-- ==== Complexity +-- +-- \(O(n)\) — exactly \(n\) multiplications and \(n\) additions +-- (or \(n - 1\) additions, depending on the fold seed). +dot :: (KnownNat n, M.Manifest r e, Num e) + => Vector n r e -> Vector n r e -> e +dot (MkVector x) (MkVector y) = + M.foldlS (+) 0 $ M.zipWith (*) x y +{-# NOINLINE [1] dot #-} + +-- | Specialised raw-array dot product for P Double. +dotP :: forall n. KnownNat n => Vector n M.P Double -> Vector n M.P Double -> Double +dotP (MkVector x) (MkVector y) = + rawDot (unwrapByteArray x) (unwrapByteArrayOffset x) + (unwrapByteArray y) (unwrapByteArrayOffset y) + (dimVal @n) +{-# NOINLINE dotP #-} + +{-# RULES "dot/P/Double" forall (x :: Vector n M.P Double) + (y :: Vector n M.P Double). + dot x y = dotP x y #-} + +-- | Scale every element of a vector by a scalar. +-- +-- __GVL4 Reference:__ Section 1.1, pp. 4–8 (scalar–vector operations). +-- +-- Computes +-- +-- \[ +-- x \;\leftarrow\; \alpha \, x +-- \] +-- +-- i.e., each component \(x_i\) is replaced by \(\alpha \, x_i\). +-- +-- ==== Type-safety guarantees +-- +-- The output vector retains the same compile-time dimension @n@ as the +-- input. +-- +-- ==== Complexity +-- +-- \(O(n)\) — one multiplication per element. +scal :: (KnownNat n, M.Manifest r e, Num e) + => e -> Vector n r e -> Vector n r e +scal alpha (MkVector x) = MkVector $ M.compute $ M.map (* alpha) x + +-- | Saxpy (Scalar Alpha X Plus Y) — the fundamental vector-update operation. +-- +-- __GVL4 Reference:__ Algorithm 1.1.1 (Saxpy), p. 4. +-- +-- Given a scalar \(\alpha\) and vectors \(x, y \in \mathbb{R}^{n}\), +-- computes +-- +-- \[ +-- y \;\leftarrow\; \alpha \, x + y +-- \] +-- +-- The Saxpy kernel is the innermost building block of the BLAS hierarchy. +-- Every Gaxpy (Level 2) and matrix–matrix (Level 3) operation can be +-- expressed as a sequence of Saxpy calls (GVL4, Section 1.1, p. 4). +-- +-- ==== Type-safety guarantees +-- +-- Both input vectors and the result share the same compile-time +-- dimension @n@. +-- +-- ==== Complexity +-- +-- \(O(n)\) — one fused multiply-add per element. +axpy :: (KnownNat n, M.Manifest r e, Num e) + => e -> Vector n r e -> Vector n r e -> Vector n r e +axpy alpha (MkVector x) (MkVector y) = + MkVector $ M.compute $ M.zipWith (\xi yi -> alpha * xi + yi) x y + +-- | Euclidean (2-) norm of a vector. +-- +-- __GVL4 Reference:__ Section 1.1, pp. 4–8 (vector norms). +-- +-- Computes +-- +-- \[ +-- \lVert x \rVert_2 \;=\; \sqrt{\sum_{i=1}^{n} x_i^{2}} +-- \] +-- +-- ==== Type-safety guarantees +-- +-- The input vector carries its length @n@ at the type level; the result +-- is a scalar of the same element type. +-- +-- ==== Complexity +-- +-- \(O(n)\) — one multiply-accumulate per element, plus a single square +-- root. +-- +-- /Note:/ This implementation does not perform the scaling trick +-- described in GVL4 (p. 5) to avoid overflow\/underflow for +-- extreme element magnitudes. For production use on +-- floating-point data with very large or very small entries, +-- consider a two-pass scaled variant. +nrm2 :: (KnownNat n, M.Manifest r e, Floating e) + => Vector n r e -> e +nrm2 (MkVector x) = sqrt $ M.foldlS (\acc xi -> acc + xi * xi) 0 x + +-- | Sum of absolute values — the 1-norm (Manhattan norm) of a vector. +-- +-- __GVL4 Reference:__ Section 1.1, pp. 4–8 (vector norms). +-- +-- Computes +-- +-- \[ +-- \lVert x \rVert_1 \;=\; \sum_{i=1}^{n} \lvert x_i \rvert +-- \] +-- +-- ==== Type-safety guarantees +-- +-- The input vector carries its length @n@ at the type level; the result +-- is a scalar of the same element type. +-- +-- ==== Complexity +-- +-- \(O(n)\) — one absolute value and one addition per element. +asum :: (KnownNat n, M.Manifest r e, Num e, Ord e) + => Vector n r e -> e +asum (MkVector x) = M.foldlS (\acc xi -> acc + abs xi) 0 x diff --git a/src/Numeric/LinearAlgebra/Massiv/BLAS/Level2.hs b/src/Numeric/LinearAlgebra/Massiv/BLAS/Level2.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/BLAS/Level2.hs @@ -0,0 +1,177 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.BLAS.Level2 +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = BLAS Level 2: Matrix–Vector Operations +-- +-- This module provides type-safe, dimension-indexed wrappers around the +-- standard BLAS Level 2 kernels. These are /matrix–vector/ operations +-- whose arithmetic cost is \(O(m \, n)\) for an \(m \times n\) matrix, +-- one level above the \(O(n)\) vector–vector operations of Level 1. +-- +-- The central operation is the /Gaxpy/ — Generalized Saxpy — which +-- computes \(y \leftarrow \alpha A x + \beta y\). It can be viewed +-- as a sequence of Saxpy updates (one per row or one per column), +-- giving rise to two natural loop orderings: +-- +-- * Golub, G. H. & Van Loan, C. F. (2013). /Matrix Computations/, +-- 4th edition (GVL4). Johns Hopkins University Press. +-- __Chapter 1, Sections 1.1.3–1.1.4__, pp. 8–12. +-- +-- Specifically: +-- +-- * __Algorithm 1.1.3__ (Row-Oriented Gaxpy, p. 8) — The @i@-th +-- component of the result is a dot product: +-- \(y_i \leftarrow a_i^{T} x + y_i\), where \(a_i^{T}\) is the +-- @i@-th row of \(A\). +-- +-- * __Algorithm 1.1.4__ (Column-Oriented Gaxpy, p. 9) — The result +-- vector is updated one column at a time via Saxpy: +-- \(y \leftarrow A(:,\!j) \, x_j + y\), for \(j = 1, \ldots, n\). +-- +-- The module also provides the rank-1 outer-product update +-- \(A \leftarrow A + \alpha x y^{T}\) (BLAS @GER@), which is the +-- matrix analogue of the Saxpy at Level 1 and plays a key role in LU +-- factorisation (GVL4, Section 3.2, p. 112). +-- +-- == Complexity +-- +-- All operations in this module are \(O(m \, n)\) for an +-- \(m \times n\) matrix. +module Numeric.LinearAlgebra.Massiv.BLAS.Level2 + ( -- * Matrix–vector multiply — Gaxpy (Algorithms 1.1.3–1.1.4, GVL4 pp. 8–9) + gemv + , matvec + , matvecP + -- * Rank-1 update (GVL4 Section 1.1.4, p. 10) + , ger + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..), Comp(..), unwrapByteArray, unwrapByteArrayOffset, + unwrapMutableByteArray, unwrapMutableByteArrayOffset) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Internal.Kernel (rawGemv) + +-- | General matrix–vector multiply (BLAS @GEMV@). +-- +-- __GVL4 Reference:__ Algorithms 1.1.3 (Row Gaxpy, p. 8) and 1.1.4 +-- (Column Gaxpy, p. 9). +-- +-- Given \(A \in \mathbb{R}^{m \times n}\), +-- \(x \in \mathbb{R}^{n}\), \(y \in \mathbb{R}^{m}\), and scalars +-- \(\alpha, \beta\), computes +-- +-- \[ +-- y \;\leftarrow\; \alpha \, A \, x \;+\; \beta \, y +-- \] +-- +-- The implementation uses a /row-oriented/ traversal (Algorithm 1.1.3): +-- for each row \(i\) the dot product \(a_i^{T} x\) is formed, scaled by +-- \(\alpha\), and added to \(\beta \, y_i\). +-- +-- ==== Type-safety guarantees +-- +-- * \(A\) is @m x n@, \(x\) is @n@, \(y\) and the result are @m@. +-- * The shared inner dimension @n@ is enforced at compile time, so a +-- dimension mismatch is a type error. +-- +-- ==== Complexity +-- +-- \(O(m \, n)\) — two floating-point operations per matrix entry +-- (one multiply, one add in the inner product), plus \(O(m)\) work +-- for the \(\alpha / \beta\) scaling. +gemv :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => e -> Matrix m n r e -> Vector n r e -> e -> Vector m r e -> Vector m r e +gemv alpha mat x beta y = + let r = dimVal @m + c = dimVal @n + in makeVector @m @r $ \i -> + let axi = foldl (\acc j -> acc + (mat ! (i, j)) * (x !. j)) 0 [0..c-1] + in alpha * axi + beta * (y !. i) + +-- | Simple matrix–vector multiply (specialisation of 'gemv'). +-- +-- __GVL4 Reference:__ Algorithm 1.1.3 (Row Gaxpy, p. 8), with +-- \(\alpha = 1\) and \(\beta = 0\). +-- +-- Given \(A \in \mathbb{R}^{m \times n}\) and +-- \(x \in \mathbb{R}^{n}\), computes +-- +-- \[ +-- y \;=\; A \, x +-- \] +-- +-- This is a convenience wrapper equivalent to @'gemv' 1 a x 0 zero@ +-- but avoids allocating or requiring an initial @y@ vector. +-- +-- ==== Type-safety guarantees +-- +-- * \(A\) is @m x n@, \(x\) is @n@, the result is @m@. +-- * The inner dimension @n@ is checked at compile time. +-- +-- ==== Complexity +-- +-- \(O(m \, n)\). +matvec :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => Matrix m n r e -> Vector n r e -> Vector m r e +matvec mat x = + let c = dimVal @n + in makeVector @m @r $ \i -> + foldl (\acc j -> acc + (mat ! (i, j)) * (x !. j)) 0 [0..c-1] +{-# NOINLINE [1] matvec #-} + +-- | Specialised raw-array matvec for P Double. +matvecP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> Vector n M.P Double -> Vector m M.P Double +matvecP (MkMatrix a) (MkVector x) = + createVector @m @M.P $ \mc -> + rawGemv (unwrapByteArray a) (unwrapByteArrayOffset a) (dimVal @n) + (unwrapByteArray x) (unwrapByteArrayOffset x) + (unwrapMutableByteArray mc) (unwrapMutableByteArrayOffset mc) + (dimVal @m) +{-# NOINLINE matvecP #-} + +{-# RULES "matvec/P/Double" forall (a :: Matrix m n M.P Double) + (x :: Vector n M.P Double). + matvec a x = matvecP a x #-} + +-- | Rank-1 update — outer product (BLAS @GER@). +-- +-- __GVL4 Reference:__ Section 1.1.4, p. 10. The rank-1 update is +-- the matrix-level analogue of the Saxpy and appears as the inner +-- kernel in outer-product formulations of LU factorisation +-- (GVL4 Section 3.2, Algorithm 3.2.1, p. 112). +-- +-- Given \(x \in \mathbb{R}^{m}\), \(y \in \mathbb{R}^{n}\), +-- \(A \in \mathbb{R}^{m \times n}\), and a scalar \(\alpha\), +-- computes +-- +-- \[ +-- A \;\leftarrow\; A \;+\; \alpha \, x \, y^{T} +-- \] +-- +-- Equivalently, each entry is updated as +-- \(a_{ij} \leftarrow a_{ij} + \alpha \, x_i \, y_j\). +-- +-- ==== Type-safety guarantees +-- +-- * \(x\) is @m@, \(y\) is @n@, \(A\) and the result are @m x n@. +-- * All dimension constraints are enforced at compile time. +-- +-- ==== Complexity +-- +-- \(O(m \, n)\) — one fused multiply-add per matrix entry. +ger :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => e -> Vector m r e -> Vector n r e -> Matrix m n r e -> Matrix m n r e +ger alpha x y mat = + makeMatrix @m @n @r $ \i j -> + (mat ! (i, j)) + alpha * (x !. i) * (y !. j) diff --git a/src/Numeric/LinearAlgebra/Massiv/BLAS/Level3.hs b/src/Numeric/LinearAlgebra/Massiv/BLAS/Level3.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/BLAS/Level3.hs @@ -0,0 +1,405 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.BLAS.Level3 +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = BLAS Level 3: Matrix–Matrix Operations +-- +-- This module provides type-safe, dimension-indexed wrappers around the +-- standard BLAS Level 3 kernels. These are /matrix–matrix/ operations +-- whose arithmetic cost is \(O(m \, n \, k)\) for an +-- \(m \times k\) by \(k \times n\) multiply, one level above the +-- \(O(m \, n)\) matrix–vector operations of Level 2. +-- +-- The algorithms implemented here correspond to the six loop orderings +-- of the triple-loop matrix multiplication described in: +-- +-- * Golub, G. H. & Van Loan, C. F. (2013). /Matrix Computations/, +-- 4th edition (GVL4). Johns Hopkins University Press. +-- __Chapter 1, Sections 1.1.5–1.1.8__, pp. 12–18. +-- +-- Specifically: +-- +-- * __Algorithm 1.1.5__ (ijk variant, p. 12) — The "row-oriented +-- inner-product" form. For each entry \(c_{ij}\) the inner product +-- of row \(i\) of \(A\) with column \(j\) of \(B\) is computed. +-- This is the variant implemented by 'gemm', 'matMul', and +-- 'matMulComp'. +-- +-- * __Algorithms 1.1.6–1.1.8__ (pp. 13–15) — The jki (Gaxpy), +-- kji (outer-product), and other orderings. These alternatives +-- differ in data-access pattern but compute the same result. The +-- present implementation uses the ijk ordering; cache-oblivious or +-- blocked variants can be added in the future. +-- +-- The module also provides elementary matrix arithmetic (addition, +-- subtraction, scaling, transpose) and a triangular matrix–matrix +-- multiply ('trmmLeft') that exploits the triangular structure to +-- halve the work. +-- +-- == Complexity +-- +-- * 'gemm', 'matMul', 'matMulComp': \(O(m \, k \, n)\). +-- * 'transpose', 'mAdd', 'mSub', 'mScale': \(O(m \, n)\). +-- * 'trmmLeft': \(O(n^{3}/2)\) (triangular, in-place structure). +module Numeric.LinearAlgebra.Massiv.BLAS.Level3 + ( -- * Matrix–matrix multiply (Algorithm 1.1.5, GVL4 p. 12) + gemm + , matMul + , matMulP + , matMulPPar + , matMulComp + -- * Elementary matrix operations (GVL4 Section 1.1, pp. 4–18) + , transpose + , transposeP + , matMulAtAP + , mAdd + , mSub + , mScale + -- * Triangular matrix multiply (GVL4 Section 1.1.8, p. 15) + , trmmLeft + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..), Comp(..), unwrapByteArray, unwrapByteArrayOffset, + unwrapMutableByteArray, unwrapMutableByteArrayOffset) +import GHC.TypeNats (KnownNat) +import GHC.Exts (Int(..), isTrue#, (>=#), (*#), (+#)) +import GHC.Prim +import GHC.ST (ST(..)) +import Data.Array.Byte (MutableByteArray(..)) +import Data.Primitive.ByteArray (ByteArray(..), newByteArray, unsafeFreezeByteArray) +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Control.Concurrent (forkOn, newEmptyMVar, putMVar, takeMVar, getNumCapabilities) +import System.IO.Unsafe (unsafePerformIO) +import GHC.IO (stToIO) +import Numeric.LinearAlgebra.Massiv.Internal.Kernel (rawGemmKernel, rawGemmBISlice, rawGemmBIBJSlice, rawSyrkLowerKernel) + +-- | Block size for cache-tiled GEMM (generic fallback path). +gemmBlockSize :: Int +gemmBlockSize = 32 +{-# INLINE gemmBlockSize #-} + +-- | General matrix multiply (BLAS @GEMM@). +-- +-- __GVL4 Reference:__ Algorithm 1.1.5 (ijk matrix multiply), p. 12. +-- +-- Given \(A \in \mathbb{R}^{m \times k}\), +-- \(B \in \mathbb{R}^{k \times n}\), +-- \(C \in \mathbb{R}^{m \times n}\), and scalars +-- \(\alpha, \beta\), computes +-- +-- \[ +-- C \;\leftarrow\; \alpha \, A \, B \;+\; \beta \, C +-- \] +gemm :: forall m k n r e. (KnownNat m, KnownNat k, KnownNat n, M.Manifest r e, Num e) + => e -> Matrix m k r e -> Matrix k n r e -> e -> Matrix m n r e -> Matrix m n r e +gemm alpha a b beta c = + let mm = dimVal @m + kk = dimVal @k + nn = dimVal @n + bs = gemmBlockSize + in createMatrix @m @n @r $ \mc -> do + -- Initialize C with β·C₀ + mapM_ (\i -> mapM_ (\j -> do + M.write_ mc (i :. j) (beta * (c ! (i, j))) + ) [0..nn-1]) [0..mm-1] + -- Tiled ikj loop: for each block triple, accumulate α·A·B + let go_bi bi = do + let iEnd = min (bi + bs) mm + let go_bk bk = do + let kEnd = min (bk + bs) kk + let go_bj bj = do + let jEnd = min (bj + bs) nn + -- Inner micro-kernel: ikj within the block + mapM_ (\i -> mapM_ (\l -> do + let aik = alpha * (a ! (i, l)) + mapM_ (\j -> do + cij <- M.readM mc (i :. j) + M.write_ mc (i :. j) (cij + aik * (b ! (l, j))) + ) [bj..jEnd-1] + ) [bk..kEnd-1]) [bi..iEnd-1] + mapM_ go_bj [0, bs .. nn-1] + mapM_ go_bk [0, bs .. kk-1] + mapM_ go_bi [0, bs .. mm-1] + +-- | Simple matrix multiply (specialisation of 'gemm'). +-- +-- __GVL4 Reference:__ Algorithm 1.1.5 (ijk matrix multiply), p. 12, +-- with \(\alpha = 1\) and \(\beta = 0\). +-- +-- Given \(A \in \mathbb{R}^{m \times k}\) and +-- \(B \in \mathbb{R}^{k \times n}\), computes +-- +-- \[ +-- C \;=\; A \, B +-- \] +matMul :: forall m k n r e. (KnownNat m, KnownNat k, KnownNat n, M.Manifest r e, Num e) + => Matrix m k r e -> Matrix k n r e -> Matrix m n r e +matMul a b = matMulGeneric a b +{-# NOINLINE [1] matMul #-} + +-- | Generic fallback for non-P or non-Double representations. +matMulGeneric :: forall m k n r e. (KnownNat m, KnownNat k, KnownNat n, M.Manifest r e, Num e) + => Matrix m k r e -> Matrix k n r e -> Matrix m n r e +matMulGeneric a b = + let mm = dimVal @m + kk = dimVal @k + nn = dimVal @n + bs = gemmBlockSize + in createMatrix @m @n @r $ \mc -> do + -- Initialize C to zero + mapM_ (\i -> mapM_ (\j -> + M.write_ mc (i :. j) 0 + ) [0..nn-1]) [0..mm-1] + -- Tiled ikj loop + let go_bi bi = do + let iEnd = min (bi + bs) mm + let go_bk bk = do + let kEnd = min (bk + bs) kk + let go_bj bj = do + let jEnd = min (bj + bs) nn + mapM_ (\i -> mapM_ (\l -> do + let aik = a ! (i, l) + mapM_ (\j -> do + cij <- M.readM mc (i :. j) + M.write_ mc (i :. j) (cij + aik * (b ! (l, j))) + ) [bj..jEnd-1] + ) [bk..kEnd-1]) [bi..iEnd-1] + mapM_ go_bj [0, bs .. nn-1] + mapM_ go_bk [0, bs .. kk-1] + mapM_ go_bi [0, bs .. mm-1] + +-- | Specialised raw-array GEMM for @P Double@. +-- Bypasses massiv's per-element abstraction and uses raw ByteArray# primops +-- with AVX2 DoubleX4# SIMD for the inner kernel. +matMulP :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) + => Matrix m k M.P Double -> Matrix k n M.P Double -> Matrix m n M.P Double +matMulP (MkMatrix arrA) (MkMatrix arrB) = + let mm = dimVal @m + kk = dimVal @k + nn = dimVal @n + baA = unwrapByteArray arrA + offA = unwrapByteArrayOffset arrA + baB = unwrapByteArray arrB + offB = unwrapByteArrayOffset arrB + in createMatrix @m @n @M.P $ \mc -> do + -- Zero-initialise C + let mbaC = unwrapMutableByteArray mc + offC = unwrapMutableByteArrayOffset mc + !(I# mm#) = mm + !(I# nn#) = nn + ST $ \s0 -> + let go i s + | isTrue# (i >=# (mm# *# nn#)) = s + | otherwise = case writeDoubleArray# (unMBA# mbaC) (unI offC +# i) 0.0## s of + s' -> go (i +# 1#) s' + in (# go 0# s0, () #) + -- Run the raw SIMD kernel + rawGemmKernel baA offA baB offB mbaC offC mm kk nn +{-# NOINLINE matMulP #-} + +-- | Parallel specialised GEMM for @P Double@. +-- Uses 2D grid partitioning when numThreads >= 4 and both dimensions are large +-- enough (min(m,n) >= 128), otherwise falls back to 1D row partitioning. +-- 2D partitioning reduces per-thread B cache traffic by a factor of sqrt(p). +-- Falls back to single-threaded 'matMulP' when @getNumCapabilities == 1@. +matMulPPar :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) + => Matrix m k M.P Double -> Matrix k n M.P Double -> Matrix m n M.P Double +matMulPPar a b = unsafePerformIO $ do + let !mm = dimVal @m + !kk = dimVal @k + !nn = dimVal @n + !baA = unwrapByteArray (unMatrix a) + !offA = unwrapByteArrayOffset (unMatrix a) + !baB = unwrapByteArray (unMatrix b) + !offB = unwrapByteArrayOffset (unMatrix b) + caps <- getNumCapabilities + -- Adaptive thread count: ensure each thread gets enough rows to + -- amortize fork/join overhead (minimum 16 rows per thread). + let !minRowsPerThread = 16 + !maxThreads = max 1 (mm `div` minRowsPerThread) + !numThreads = min caps (min mm maxThreads) + if numThreads <= 1 + then pure (matMulP a b) + else do + -- Allocate mutable C, zero-initialise + mc <- stToIO $ M.newMArray (Sz (mm :. nn)) (0 :: Double) + let !mbaC = unwrapMutableByteArray mc + !offC = unwrapMutableByteArrayOffset mc + -- Choose 1D or 2D decomposition + let !use2D = numThreads >= 4 && mm >= 128 && nn >= 128 + if use2D + then do + -- 2D grid: pr rows × pc columns, pr * pc = numThreads + -- Choose pr, pc to balance aspect ratio: pr/pc ≈ mm/nn + let (pr, pc) = gridDims numThreads mm nn + !rChunk = (mm + pr - 1) `div` pr + !cChunk = (nn + pc - 1) `div` pc + mvars <- sequence + [ do let !biStart = tr * rChunk + !biEnd = min (biStart + rChunk) mm + !bjStart = tc * cChunk + !bjEnd = min (bjStart + cChunk) nn + mv <- newEmptyMVar + _ <- forkOn (tr * pc + tc) $ do + stToIO $ rawGemmBIBJSlice baA offA baB offB mbaC offC + biStart biEnd bjStart bjEnd mm kk nn + putMVar mv () + pure mv + | tr <- [0..pr-1], tc <- [0..pc-1] + ] + mapM_ takeMVar mvars + else do + -- 1D row partitioning (original path) + let !chunkSize = (mm + numThreads - 1) `div` numThreads + mvars <- mapM (\t -> do + let !biStart = t * chunkSize + !biEnd = min (biStart + chunkSize) mm + mv <- newEmptyMVar + _ <- forkOn t $ do + stToIO $ rawGemmBISlice baA offA baB offB mbaC offC biStart biEnd mm kk nn + putMVar mv () + pure mv + ) [0..numThreads-1] + mapM_ takeMVar mvars + -- Freeze and wrap + arr <- stToIO $ M.freezeS mc + pure (MkMatrix arr) +{-# NOINLINE matMulPPar #-} + +-- | Compute 2D grid dimensions (pr × pc) for p threads such that +-- pr * pc = p and the aspect ratio pr/pc approximates m/n. +gridDims :: Int -> Int -> Int -> (Int, Int) +gridDims p m n = + let sqrtP = floor (sqrt (fromIntegral p :: Double)) :: Int + -- Try all factorizations of p and pick the one with best aspect match + factors = [(i, p `div` i) | i <- [1..sqrtP], p `mod` i == 0] + targetRatio = fromIntegral m / fromIntegral (max 1 n) :: Double + score (pr, pc) = abs (fromIntegral pr / fromIntegral (max 1 pc) - targetRatio) + best = minimumBy (\a' b' -> compare (score a') (score b')) factors + -- Also consider the transpose (pc, pr) + bestT = let (pr, pc) = best in if score (pc, pr) < score best then (pc, pr) else best + in bestT + where + minimumBy _ [x] = x + minimumBy f (x:xs) = foldl' (\a' b' -> if f a' b' == GT then b' else a') x xs + minimumBy _ [] = (1, p) -- fallback + +{-# RULES +"matMul/P/Double" forall (a :: Matrix m k M.P Double) + (b :: Matrix k n M.P Double). + matMul a b = matMulP a b + #-} + +-- | Matrix multiply with explicit computation strategy. +matMulComp :: forall m k n r e. (KnownNat m, KnownNat k, KnownNat n, M.Manifest r e, Num e) + => Comp -> Matrix m k r e -> Matrix k n r e -> Matrix m n r e +matMulComp comp a b = + case comp of + Seq -> matMul a b + _ -> -- For parallel: use delayed array with ikj-reordered inner product + let kk = dimVal @k + in makeMatrixComp @m @n @r comp $ \i j -> + foldl' (\acc l -> acc + (a ! (i, l)) * (b ! (l, j))) 0 [0..kk-1] + +-- | Matrix transpose. +transpose :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e) + => Matrix m n r e -> Matrix n m r e +transpose (MkMatrix arr) = MkMatrix $ M.compute $ M.transposeInner arr + +-- | P-specialised raw-primop matrix transpose. +-- Avoids per-element overhead of massiv's delayed transpose. +transposeP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> Matrix n m M.P Double +transposeP (MkMatrix a) = + let !mm = dimVal @m + !nn = dimVal @n + !(I# mm#) = mm + !(I# nn#) = nn + !(ByteArray ba#) = unwrapByteArray a + !(I# off#) = unwrapByteArrayOffset a + in createMatrix @n @m @M.P $ \mu -> + let !(MutableByteArray mba#) = unwrapMutableByteArray mu + !(I# offR#) = unwrapMutableByteArrayOffset mu + in ST $ \s0 -> + -- Iterate source rows (sequential read, strided write) + let goRow j s + | isTrue# (j >=# mm#) = s + | otherwise = + let goCol i s1 + | isTrue# (i >=# nn#) = s1 + | otherwise = + let !src = off# +# j *# nn# +# i + !dst = offR# +# i *# mm# +# j + !v = indexDoubleArray# ba# src + in case writeDoubleArray# mba# dst v s1 of + s2 -> goCol (i +# 1#) s2 + in goRow (j +# 1#) (goCol 0# s) + in (# goRow 0# s0, () #) +{-# NOINLINE transposeP #-} + +-- | P-specialised A^T * A without materialising A^T. +-- Computes C = A^T * A using a fused DSYRK kernel that processes only the +-- lower triangle (halving flops) and mirrors to the upper triangle. +-- Avoids materialising A^T entirely — one allocation, no transpose pass. +matMulAtAP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> Matrix n n M.P Double +matMulAtAP (MkMatrix arrA) = + let mm = dimVal @m + nn = dimVal @n + baA = unwrapByteArray arrA + offA = unwrapByteArrayOffset arrA + in createMatrix @n @n @M.P $ \mc -> do + -- Zero-initialise C (n×n) + let mbaC = unwrapMutableByteArray mc + offC = unwrapMutableByteArrayOffset mc + !(I# nn#) = nn + ST $ \s0 -> + let go i s + | isTrue# (i >=# (nn# *# nn#)) = s + | otherwise = case writeDoubleArray# (unMBA# mbaC) (unI offC +# i) 0.0## s of + s' -> go (i +# 1#) s' + in (# go 0# s0, () #) + -- Run the fused SYRK kernel: C = A^T * A (lower triangle + mirror) + rawSyrkLowerKernel baA offA mbaC offC mm nn +{-# NOINLINE matMulAtAP #-} + +-- | Element-wise matrix addition. +mAdd :: (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => Matrix m n r e -> Matrix m n r e -> Matrix m n r e +mAdd (MkMatrix a) (MkMatrix b) = MkMatrix $ M.compute $ M.zipWith (+) a b + +-- | Element-wise matrix subtraction. +mSub :: (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => Matrix m n r e -> Matrix m n r e -> Matrix m n r e +mSub (MkMatrix a) (MkMatrix b) = MkMatrix $ M.compute $ M.zipWith (-) a b + +-- | Scalar–matrix multiply. +mScale :: (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => e -> Matrix m n r e -> Matrix m n r e +mScale alpha (MkMatrix a) = MkMatrix $ M.compute $ M.map (* alpha) a + +-- | Left-multiply by a lower-triangular matrix (BLAS @TRMM@, left side). +trmmLeft :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Matrix n n r e -> Matrix n n r e -> Matrix n n r e +trmmLeft l b = + let nn = dimVal @n + in makeMatrix @n @n @r $ \i j -> + foldl' (\acc k -> acc + (l ! (i, k)) * (b ! (k, j))) 0 [0..min i (nn-1)] + +-- Helpers to unwrap newtypes for raw primop access +unMBA# :: MutableByteArray s -> MutableByteArray# s +unMBA# (MutableByteArray mba) = mba +{-# INLINE unMBA# #-} + +unI :: Int -> Int# +unI (I# i) = i +{-# INLINE unI #-} diff --git a/src/Numeric/LinearAlgebra/Massiv/Eigen/Hessenberg.hs b/src/Numeric/LinearAlgebra/Massiv/Eigen/Hessenberg.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Eigen/Hessenberg.hs @@ -0,0 +1,127 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Reduction of a general square matrix to upper Hessenberg form via +-- Householder similarity transformations, following Golub & Van Loan, +-- /Matrix Computations/, 4th edition (GVL4), Section 7.4, pp. 383--393. +-- +-- An upper Hessenberg matrix \(H\) satisfies \(h_{ij} = 0\) for +-- \(i > j + 1\); that is, all entries below the first subdiagonal are zero. +-- Given \(A \in \mathbb{R}^{n \times n}\), the algorithm computes an +-- orthogonal matrix \(Q\) (accumulated as a product of \(n - 2\) Householder +-- reflectors) such that +-- +-- \[ +-- A = Q H Q^T +-- \] +-- +-- This is the standard first phase in eigenvalue algorithms (e.g. the +-- implicit QR algorithm in "Numeric.LinearAlgebra.Massiv.Eigen.Schur"), +-- because QR steps preserve Hessenberg structure and are much cheaper on a +-- Hessenberg matrix than on a full matrix. +-- +-- __Algorithm:__ Householder reduction to Hessenberg form (GVL4 +-- Algorithm 7.4.2, p. 387). +-- +-- __Complexity:__ \(\frac{10}{3} n^3\) floating-point operations. +module Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg + ( -- * Hessenberg reduction (Algorithm 7.4.2) + hessenberg + , hessenbergInPlace + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal + +-- | Reduce a general matrix to upper Hessenberg form (GVL4 Algorithm 7.4.2, +-- p. 387). +-- +-- Computes the factorisation +-- +-- \[ +-- A = Q \, H \, Q^T +-- \] +-- +-- where \(Q\) is orthogonal (a product of \(n - 2\) Householder reflectors) +-- and \(H\) is upper Hessenberg, i.e. \(h_{ij} = 0\) for \(i > j + 1\). +-- +-- At step \(k\) the algorithm determines a Householder reflector +-- \(P_k = I - \beta_k v_k v_k^T\) that zeroes entries \(k+2, \ldots, n\) in +-- column \(k\) of the current matrix, and applies it as a similarity +-- transformation \(H \leftarrow P_k H P_k\). +-- +-- __Complexity:__ \(\frac{10}{3} n^3\) flops (GVL4, p. 389). +-- +-- Returns @(Q, H)@. +hessenberg :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> (Matrix n n r e, Matrix n n r e) +hessenberg a = + let nn = dimVal @n + go :: Int -> Matrix n n r e -> Matrix n n r e -> (Matrix n n r e, Matrix n n r e) + go k q_ h_ + | k >= nn - 2 = (q_, h_) + | otherwise = + -- Compute Householder to zero out h(k+2:n, k) + let x0 = h_ ! (k+1, k) + sigma = foldl' (\acc i -> acc + (h_ ! (i, k)) * (h_ ! (i, k))) 0 [k+2..nn-1] + in if sigma == 0 && x0 >= 0 + then go (k + 1) q_ h_ + else + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + -- v is zero for indices 0..k, v(k+1)=1, v(i) = h(i,k)/v0 for i > k+1 + v = makeVector @n @r $ \i -> + if i <= k then 0 + else if i == k + 1 then 1 + else (h_ ! (i, k)) / v0 + -- H ← (I - β·v·vᵀ)·H·(I - β·v·vᵀ) = similarity transform + -- First: H ← (I - β·v·vᵀ)·H + h1 = applyFromLeft v beta h_ + -- Then: H ← H·(I - β·v·vᵀ) + h2 = applyFromRight h1 v beta + -- Q ← Q·(I - β·v·vᵀ) + q_new = applyFromRight q_ v beta + in go (k + 1) q_new h2 + + q0 = identityMatrix @n @r + in go 0 q0 a + +-- Helper: apply Householder from left +applyFromLeft :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Vector n r e -> e -> Matrix n n r e -> Matrix n n r e +applyFromLeft v beta h = + let nn = dimVal @n + in makeMatrix @n @n @r $ \i j -> + let wj = beta * foldl' (\acc k -> acc + (v !. k) * (h ! (k, j))) 0 [0..nn-1] + in (h ! (i, j)) - (v !. i) * wj + +-- Helper: apply Householder from right +applyFromRight :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Matrix n n r e -> Vector n r e -> e -> Matrix n n r e +applyFromRight h v beta = + let nn = dimVal @n + in makeMatrix @n @n @r $ \i j -> + let wi = beta * foldl' (\acc k -> acc + (h ! (i, k)) * (v !. k)) 0 [0..nn-1] + in (h ! (i, j)) - wi * (v !. j) + +-- | Compute only the upper Hessenberg matrix \(H\), discarding the +-- orthogonal factor \(Q\). +-- +-- This is a convenience wrapper around 'hessenberg' that returns only the +-- second component of the pair. Use this when only the Hessenberg form is +-- needed and the transformation matrix is not required (e.g. when computing +-- eigenvalues but not eigenvectors). +hessenbergInPlace :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Matrix n n r e +hessenbergInPlace a = snd (hessenberg a) diff --git a/src/Numeric/LinearAlgebra/Massiv/Eigen/Power.hs b/src/Numeric/LinearAlgebra/Massiv/Eigen/Power.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Eigen/Power.hs @@ -0,0 +1,182 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Eigen.Power +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Power iteration methods for computing individual eigenvalue\/eigenvector +-- pairs of a general square matrix. +-- +-- This module implements three iterative projection techniques drawn from +-- Golub & Van Loan, /Matrix Computations/, 4th edition (GVL4), Section 7.3, +-- pp. 372--382: +-- +-- * __Power method__ (Algorithm 7.3.3, p. 375) — converges to the dominant +-- eigenpair at a rate governed by the ratio \(|\lambda_2 / \lambda_1|\) per +-- iteration, where \(\lambda_1\) is the eigenvalue of largest modulus. +-- +-- * __Inverse iteration__ (Section 7.3.1, p. 377) — given a shift \(\mu\), +-- converges to the eigenvalue closest to \(\mu\) by applying the power +-- method to \((A - \mu I)^{-1}\). +-- +-- * __Rayleigh quotient iteration__ (Section 7.3.2, p. 379) — an adaptive +-- variant of inverse iteration in which the shift is updated at every step +-- to equal the current Rayleigh quotient. For symmetric matrices this +-- achieves /cubic/ convergence; for general matrices the convergence is +-- /quadratic/. +-- +-- All three routines return an approximate eigenvalue \(\lambda\) and its +-- associated eigenvector \(q\) satisfying \(Aq \approx \lambda q\). +module Numeric.LinearAlgebra.Massiv.Eigen.Power + ( -- * Power method (Algorithm 7.3.3) + powerMethod + -- * Inverse iteration (Section 7.3.1) + , inverseIteration + -- * Rayleigh quotient iteration (Section 7.3.2) + , rayleighQuotient + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level1 (dot, scal, nrm2) +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvec) +import Numeric.LinearAlgebra.Massiv.Solve.LU (luSolve) + +-- | Power method for computing the dominant eigenpair (GVL4 Algorithm 7.3.3, +-- p. 375). +-- +-- Given a square matrix \(A \in \mathbb{R}^{n \times n}\) with eigenvalues +-- ordered \(|\lambda_1| > |\lambda_2| \geq \cdots \geq |\lambda_n|\), the +-- power method generates a sequence of vectors +-- +-- \[ +-- z_k = A q_{k-1}, \qquad q_k = z_k / \|z_k\|_2 +-- \] +-- +-- that converges to the eigenvector associated with \(\lambda_1\). The +-- corresponding eigenvalue is estimated via the Rayleigh quotient +-- \(\lambda \approx q_k^T A q_k\). +-- +-- __Convergence rate:__ the error contracts by a factor of +-- \(|\lambda_2 / \lambda_1|\) per iteration (GVL4, p. 375). The method +-- therefore requires a dominant eigenvalue that is well-separated from the +-- rest of the spectrum. +-- +-- Returns @(eigenvalue, eigenvector)@ once the eigenvalue estimate changes by +-- less than the given tolerance, or after the specified number of iterations. +powerMethod :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e + -> Vector n r e -- ^ Initial guess \(q_0\) (should be unit norm) + -> Int -- ^ Maximum iterations + -> e -- ^ Convergence tolerance + -> (e, Vector n r e) +powerMethod a q0 maxIter tol = go 0 q0 0 + where + go :: Int -> Vector n r e -> e -> (e, Vector n r e) + go iter q prevLambda + | iter >= maxIter = (prevLambda, q) + | otherwise = + let z = matvec a q -- z = A·q + znorm = nrm2 z -- ‖z‖₂ + qNew = scal (1 / znorm) z -- q = z / ‖z‖₂ + lambda = dot qNew (matvec a qNew) -- λ = qᵀAq (Rayleigh quotient) + in if abs (lambda - prevLambda) < tol + then (lambda, qNew) + else go (iter + 1) qNew lambda + +-- | Inverse iteration for computing the eigenpair closest to a given shift +-- (GVL4 Section 7.3.1, p. 377). +-- +-- Given a shift \(\mu\) that approximates some eigenvalue of \(A\), inverse +-- iteration applies the power method to the matrix \((A - \mu I)^{-1}\). +-- Each step solves the linear system +-- +-- \[ +-- (A - \mu I)\, z_k = q_{k-1}, \qquad q_k = z_k / \|z_k\|_2 +-- \] +-- +-- and converges to the eigenvalue \(\lambda_j\) that minimises +-- \(|\lambda_j - \mu|\). The convergence rate is +-- \(|\lambda_j - \mu| / |\lambda_i - \mu|\) per iteration, where +-- \(\lambda_i\) is the second-closest eigenvalue to \(\mu\). +-- +-- The eigenvalue estimate is refined at each step via the Rayleigh quotient +-- \(\lambda \approx q_k^T A q_k\) (using the /original/ matrix \(A\)). +-- +-- Returns @(eigenvalue, eigenvector)@. +inverseIteration :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e + -> e -- ^ Shift \(\mu\) + -> Vector n r e -- ^ Initial guess + -> Int -- ^ Maximum iterations + -> e -- ^ Convergence tolerance + -> (e, Vector n r e) +inverseIteration a mu q0 maxIter tol = go 0 q0 0 + where + aShifted = makeMatrix @n @n @r $ \i j -> + if i == j then (a ! (i, j)) - mu else a ! (i, j) + + go :: Int -> Vector n r e -> e -> (e, Vector n r e) + go iter q prevLambda + | iter >= maxIter = (prevLambda, q) + | otherwise = + let z = luSolve aShifted q -- Solve (A - μI)z = q + znorm = nrm2 z + qNew = scal (1 / znorm) z + lambda = dot qNew (matvec a qNew) -- Rayleigh quotient with original A + in if abs (lambda - prevLambda) < tol + then (lambda, qNew) + else go (iter + 1) qNew lambda + +-- | Rayleigh quotient iteration (GVL4 Section 7.3.2, p. 379). +-- +-- An adaptive variant of inverse iteration in which the shift \(\mu_k\) is +-- set equal to the current Rayleigh quotient at every step: +-- +-- \[ +-- \mu_k = q_k^T A q_k, \qquad (A - \mu_k I)\, z_{k+1} = q_k, \qquad +-- q_{k+1} = z_{k+1} / \|z_{k+1}\|_2 +-- \] +-- +-- __Convergence:__ +-- +-- * For /symmetric/ matrices \(A = A^T\), the iteration converges +-- /cubically/ — the residual \(\|Aq - \lambda q\|\) is cubed at each step +-- (GVL4, p. 379). +-- * For general (non-symmetric) matrices, convergence is /quadratic/. +-- +-- Because the shift changes at each iteration, a fresh LU factorisation of +-- \(A - \mu_k I\) is computed every step. Despite this extra cost the rapid +-- convergence usually makes Rayleigh quotient iteration the method of choice +-- when a good initial vector is available. +-- +-- Returns @(eigenvalue, eigenvector)@. +rayleighQuotient :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e + -> Vector n r e -- ^ Initial guess + -> Int -- ^ Maximum iterations + -> e -- ^ Convergence tolerance + -> (e, Vector n r e) +rayleighQuotient a q0 maxIter tol = go 0 q0 (dot q0 (matvec a q0)) + where + go :: Int -> Vector n r e -> e -> (e, Vector n r e) + go iter q lambda + | iter >= maxIter = (lambda, q) + | otherwise = + let nn = dimVal @n + aShifted = makeMatrix @n @n @r $ \i j -> + if i == j then (a ! (i, j)) - lambda else a ! (i, j) + z = luSolve aShifted q + znorm = nrm2 z + qNew = scal (1 / znorm) z + lambdaNew = dot qNew (matvec a qNew) + in if abs (lambdaNew - lambda) < tol + then (lambdaNew, qNew) + else go (iter + 1) qNew lambdaNew diff --git a/src/Numeric/LinearAlgebra/Massiv/Eigen/SVD.hs b/src/Numeric/LinearAlgebra/Massiv/Eigen/SVD.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Eigen/SVD.hs @@ -0,0 +1,1992 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UnboxedTuples #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Eigen.SVD +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Singular Value Decomposition (SVD) of a general real matrix, following +-- Golub & Van Loan, /Matrix Computations/, 4th edition (GVL4), Section 8.6, +-- pp. 498--512. +-- +-- __Theorem 8.6.1 (SVD Existence, p. 499):__ For any +-- \(A \in \mathbb{R}^{m \times n}\) with \(m \geq n\) there exist orthogonal +-- matrices \(U \in \mathbb{R}^{m \times m}\) and +-- \(V \in \mathbb{R}^{n \times n}\) such that +-- +-- \[ +-- A = U \, \Sigma \, V^T, \qquad +-- \Sigma = \mathrm{diag}(\sigma_1, \ldots, \sigma_n), +-- \qquad \sigma_1 \geq \sigma_2 \geq \cdots \geq \sigma_n \geq 0 +-- \] +-- +-- The \(\sigma_i\) are the /singular values/ of \(A\) and equal the +-- non-negative square roots of the eigenvalues of \(A^T A\). +module Numeric.LinearAlgebra.Massiv.Eigen.SVD + ( -- * Full SVD + svd + , svdP + , svdAtAP + , svdGKP + -- * Singular values only + , singularValues + , singularValuesP + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..), unwrapByteArray, unwrapByteArrayOffset, + unwrapMutableByteArray, unwrapMutableByteArrayOffset) +import Data.Primitive.ByteArray (ByteArray(..), MutableByteArray(..), newByteArray, + unsafeFreezeByteArray) +import GHC.TypeNats (KnownNat) +import Control.Monad (forM_, when) +import Control.Monad.ST (runST) +import Data.List (sortBy) +import Data.Ord () +import GHC.Exts +import GHC.ST (ST(..)) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, transpose, matMulAtAP) +-- matvecP no longer needed: U-matrix now computed via single GEMM +import Numeric.LinearAlgebra.Massiv.Eigen.Symmetric + ( symmetricEigen, symmetricEigenP, symmetricEigenPDC + -- D&C secular equation infrastructure (reused for bidiagonal SVD) + , secularSolve, deflatePartition, dcEigenvectors + , sumZSq + , readRawI + ) +import Numeric.LinearAlgebra.Massiv.Internal.Kernel + ( rawMutSumSqColumn, rawMutSumSqRow + , rawMutHouseholderApply, rawMutHouseholderApplyRow + , rawMutQAccum + , rawMutApplyGivensColumns + , rawMutApplyGivensColumnsCM + , rawTransposeToColMajor, rawTransposeFromColMajor + , rawGemmKernel, rawZeroDoubles, rawNegateDoubles + , rawCopyColumn ) + +-- | Compute the full Singular Value Decomposition (GVL4 Theorem 8.6.1, +-- p. 499). +-- +-- For an \(m \times n\) matrix \(A\) with \(m \geq n\), computes +-- +-- \[ +-- A = U \, \Sigma \, V^T +-- \] +-- +-- where +-- +-- * \(U \in \mathbb{R}^{m \times m}\) is orthogonal (columns are the +-- /left singular vectors/), +-- * \(\Sigma = \mathrm{diag}(\sigma_1, \ldots, \sigma_n)\) with +-- \(\sigma_1 \geq \cdots \geq \sigma_n \geq 0\) (the /singular values/), +-- * \(V \in \mathbb{R}^{n \times n}\) is orthogonal (columns are the +-- /right singular vectors/). +-- +-- __Method:__ Forms \(A^T A\) and calls +-- 'Numeric.LinearAlgebra.Massiv.Eigen.Symmetric.symmetricEigen' to obtain +-- the eigendecomposition \(A^T A = V \Lambda V^T\). Singular values are +-- recovered as \(\sigma_i = \sqrt{\max(0, \lambda_i)}\) and left singular +-- vectors as \(u_i = A v_i / \sigma_i\). For zero singular values the +-- corresponding column of \(U\) is set to the appropriate standard basis +-- vector. +-- +-- Returns @(U, sigma, V)@. +svd :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix m n r e + -> (Matrix m m r e, Vector n r e, Matrix n n r e) +svd a = + let nn = dimVal @n + at = transpose a + ata = matMul at a -- n×n symmetric positive semidefinite + -- Eigendecomposition of AᵀA + (eigvalsRaw_, vRaw_) = symmetricEigen ata (30 * nn) 1e-12 + -- Sort eigenvalues descending; build O(1) permutation array + permBA_ = buildPermArray + (map snd $ sortBy (\(a_,_) (b_,_) -> compare (Down a_) (Down b_)) + [(eigvalsRaw_ !. i, i) | i <- [0..nn-1]]) + nn + v = makeMatrix @n @n @r $ \i j -> vRaw_ ! (i, indexPermArray permBA_ j) + -- Singular values = sqrt of sorted eigenvalues (clamp negatives to 0) + sigma = makeVector @n @r $ \j -> + let ev = eigvalsRaw_ !. indexPermArray permBA_ j + in if ev > 0 then sqrt ev else 0 + -- Compute U: u_i = A·v_i / σ_i + -- First, build U by computing A·V column by column + u = makeMatrix @m @m @r $ \i j -> + if j < nn then + let sj = sigma !. j + in if sj > 1e-14 + then -- u_j = (1/σ_j) · Σ_k A(i,k) · V(k,j) + let av_ = foldl' (\acc k -> acc + (a ! (i, k)) * (v ! (k, j))) 0 [0..nn-1] + in av_ / sj + else -- Zero singular value; use arbitrary orthogonal vector + if i == j then 1 else 0 + else + -- Extra columns for m > n: extend to full orthogonal basis + if i == j then 1 else 0 + in (u, sigma, v) + +-- | P-specialised full SVD using raw ByteArray# SIMD kernels throughout. +-- +-- Wires 'matMulP' (SIMD GEMM), 'symmetricEigenP' (raw primop QR iteration), +-- and 'matvecP' (SIMD matrix–vector product) into the SVD pipeline. +-- | P-specialised full SVD. Uses the A^T A eigendecomposition path by default +-- as it is currently faster than the Golub-Kahan bidiagonalisation path +-- (svdGKP) at all sizes. The GK path will become the default once blocked +-- bidiagonalisation is implemented. +svdP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double + -> (Matrix m m M.P Double, Vector n M.P Double, Matrix n n M.P Double) +svdP = svdAtAP +{-# NOINLINE svdP #-} + +-- | SVD via A^T A eigendecomposition. +-- Forms A^T A, eigendecomposes via 'symmetricEigenP', recovers singular +-- values as square roots and left singular vectors via matrix-vector products. +svdAtAP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double + -> (Matrix m m M.P Double, Vector n M.P Double, Matrix n n M.P Double) +svdAtAP a = + let !mm = dimVal @m + !nn = dimVal @n + !ata = matMulAtAP a -- n×n symmetric positive semidefinite, fast transpose + SIMD GEMM + -- Eigendecomposition of AᵀA: D&C for large, QR iteration for small + (!eigvalsRaw, !vRaw) = if nn >= 50 + then symmetricEigenPDC ata 1e-12 + else symmetricEigenP ata (max 30 (6 * nn)) 1e-12 + -- Sort eigenvalues descending; build O(1) permutation array + !permList = map snd $ sortBy (\(a_,_) (b_,_) -> compare (Down a_) (Down b_)) + [(eigvalsRaw !. i, i) | i <- [0..nn-1]] + !permBA = buildPermArray permList nn + -- Rearrange V columns via O(1) indexed permutation using raw ByteArray copy + !v = createMatrix @n @n @M.P $ \mv -> do + let !baV = unwrapByteArray (unMatrix vRaw) + !offV = unwrapByteArrayOffset (unMatrix vRaw) + !mbaVP = unwrapMutableByteArray mv + !offVP = unwrapMutableByteArrayOffset mv + !(ByteArray baV#) = baV + !(I# offV#) = offV + !(MutableByteArray mbaVP#) = mbaVP + !(I# offVP#) = offVP + !(I# nnV#) = nn + -- Copy columns: V_new[i,j] = V_raw[i, perm[j]] + ST $ \s0 -> + let goRow i s + | isTrue# (i >=# nnV#) = s + | otherwise = + let goCol j s1 + | isTrue# (j >=# nnV#) = s1 + | otherwise = + let !(I# pj) = indexPermArray permBA (I# j) + !val = indexDoubleArray# baV# (offV# +# i *# nnV# +# pj) + in case writeDoubleArray# mbaVP# (offVP# +# i *# nnV# +# j) val s1 of + s2 -> goCol (j +# 1#) s2 + in goRow (i +# 1#) (goCol 0# s) + in (# goRow 0# s0, () #) + -- Singular values = sqrt of sorted eigenvalues (clamp negatives to 0) + sigma = makeVector @n @M.P $ \j -> + let !(ByteArray baEV#) = unwrapByteArray (unVector eigvalsRaw) + !(I# offEV#) = unwrapByteArrayOffset (unVector eigvalsRaw) + !(I# pj#) = indexPermArray permBA j + ev = case indexDoubleArray# baEV# (offEV# +# pj#) of v_ -> D# v_ + in if ev > 0 then sqrt ev else 0 + -- Compute U = A · V · diag(1/σ) via pre-scaled V and single GEMM. + -- This avoids the intermediate m×n AV matrix and a separate scaling pass. + -- V_scaled[i,j] = V[i,j] / σ_j (zero for σ_j ≤ ε). + !vScaled = createMatrix @n @n @M.P @Double $ \mvs -> do + let !baV = unwrapByteArray (unMatrix v) + !offVs = unwrapByteArrayOffset (unMatrix v) + !mbaVS = unwrapMutableByteArray mvs + !offVS = unwrapMutableByteArrayOffset mvs + -- Pre-compute invSigma + mbaInvS <- newByteArray (nn * 8) + forM_ [0..nn-1] $ \j -> do + let sj = sigma !. j + writeRawD mbaInvS 0 j (if sj > 1e-14 then 1.0 / sj else 0.0) + !(ByteArray baInvS#) <- unsafeFreezeByteArray mbaInvS + -- Scale each column: V_scaled[i,j] = V[i,j] * invSigma[j] + let !(ByteArray baV#) = baV + !(I# offVs#) = offVs + !(MutableByteArray mbaVS#) = mbaVS + !(I# offVS#) = offVS + !(I# nnV#) = nn + !nn4 = nn - (nn `rem` 4) + !(I# nn4#) = nn4 + ST $ \s0 -> + let goRow i s + | isTrue# (i >=# nnV#) = s + | otherwise = + let !srcOff = offVs# +# i *# nnV# + !dstOff = offVS# +# i *# nnV# + goSimd j s1 + | isTrue# (j >=# nn4#) = s1 + | otherwise = + let vv = indexDoubleArrayAsDoubleX4# baV# (srcOff +# j) + sv = indexDoubleArrayAsDoubleX4# baInvS# j + !p = timesDoubleX4# vv sv + in case writeDoubleArrayAsDoubleX4# mbaVS# (dstOff +# j) p s1 of + s2 -> goSimd (j +# 4#) s2 + goScalar j s1 + | isTrue# (j >=# nnV#) = s1 + | otherwise = + let vVal = indexDoubleArray# baV# (srcOff +# j) + sVal = indexDoubleArray# baInvS# j + in case writeDoubleArray# mbaVS# (dstOff +# j) (vVal *## sVal) s1 of + s2 -> goScalar (j +# 1#) s2 + in goRow (i +# 1#) (goScalar nn4# (goSimd 0# s)) + in (# goRow 0# s0, () #) + -- U = A · V_scaled: GEMM writes m×n result directly. + -- For mm == nn (square), GEMM writes directly into U. + -- For mm > nn (rectangular), GEMM writes into temp then copy columns. + u = createMatrix @m @m @M.P $ \mu -> do + let !mbaU = unwrapMutableByteArray mu + !offU = unwrapMutableByteArrayOffset mu + !(I# mm#) = mm + -- Zero all of U + rawZeroDoubles mbaU offU (mm * mm) + let !baA = unwrapByteArray (unMatrix a) + !offA = unwrapByteArrayOffset (unMatrix a) + !baVS = unwrapByteArray (unMatrix vScaled) + !offVS = unwrapByteArrayOffset (unMatrix vScaled) + if mm == nn + then + -- Direct GEMM into U (stride mm == nn, so layout matches) + rawGemmKernel baA offA baVS offVS mbaU offU mm nn nn + else do + -- GEMM into temp (m×n), then copy columns into U (m×m) + mbaTemp <- newByteArray (mm * nn * 8) + rawZeroDoubles mbaTemp 0 (mm * nn) + rawGemmKernel baA offA baVS offVS mbaTemp 0 mm nn nn + baTemp <- unsafeFreezeByteArray mbaTemp + -- Copy: U[i, 0..nn-1] = temp[i, 0..nn-1] + let !(ByteArray baT#) = baTemp + !(MutableByteArray mbaU#) = mbaU + !(I# offU#) = offU + !(I# nn#) = nn + ST $ \s0 -> + let goCopy i s + | isTrue# (i >=# mm#) = s + | otherwise = + let goCol j s1 + | isTrue# (j >=# nn#) = s1 + | otherwise = + let !val = indexDoubleArray# baT# (i *# nn# +# j) + in case writeDoubleArray# mbaU# (offU# +# i *# mm# +# j) val s1 of + s2 -> goCol (j +# 1#) s2 + in goCopy (i +# 1#) (goCol 0# s) + in (# goCopy 0# s0, () #) + -- Fix zero singular values: set diagonal U[j,j] = 1.0 + forM_ [0..nn-1] $ \j -> do + let sj = sigma !. j + when (sj <= 1e-14) $ + writeRawD mbaU offU (j * mm + j) 1.0 + -- Extra columns for m > n: extend to full orthogonal basis + forM_ [nn..mm-1] $ \j -> + writeRawD mbaU offU (j * mm + j) 1.0 + in (u, sigma, v) +{-# NOINLINE svdAtAP #-} + +-- | Read a Double from an immutable ByteArray at element index. +readBA :: ByteArray -> Int -> Int -> Double +readBA (ByteArray ba) (I# off) (I# i) = + case indexDoubleArray# ba (off +# i) of v -> D# v +{-# INLINE readBA #-} + +-- | Build an unboxed Int permutation array from a list for O(1) indexed access. +buildPermArray :: [Int] -> Int -> ByteArray +buildPermArray xs n = runST $ do + mba <- newByteArray (n * 8) -- 8 bytes per Int on 64-bit + let go _ [] = pure () + go i (x:rest) = do + let !(MutableByteArray mba#) = mba + !(I# i#) = i + !(I# x#) = x + ST $ \s -> case writeIntArray# mba# i# x# s of s' -> (# s', () #) + go (i + 1) rest + go 0 xs + unsafeFreezeByteArray mba +{-# INLINE buildPermArray #-} + +-- | O(1) index into a permutation ByteArray. +indexPermArray :: ByteArray -> Int -> Int +indexPermArray (ByteArray ba#) (I# i#) = + case indexIntArray# ba# i# of x# -> I# x# +{-# INLINE indexPermArray #-} + +-- | Compute only the singular values of \(A\), sorted in descending order. +singularValues :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix m n r e -> Vector n r e +singularValues a = + let nn = dimVal @n + at = transpose a + ata = matMul at a + (eigvals, _) = symmetricEigen ata (30 * nn) 1e-12 + -- Sort eigenvalues descending, take sqrt + evList = map (\i -> eigvals !. i) [0..nn-1] + sorted = sortBy (\x y -> compare (Down x) (Down y)) evList + in makeVector @n @r $ \i -> + let ev = sorted !! i + in if ev > 0 then sqrt ev else 0 + +-- | P-specialised singular values using raw SIMD GEMM and raw primop eigenvalue solver. +singularValuesP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> Vector n M.P Double +singularValuesP a = + let nn = dimVal @n + !ata = matMulAtAP a + (!eigvals, _) = symmetricEigenP ata (10 * nn) 1e-12 + evList = map (\i -> eigvals !. i) [0..nn-1] + sorted = sortBy (\x y -> compare (Down x) (Down y)) evList + in makeVector @n @M.P $ \i -> + let ev = sorted !! i + in if ev > 0 then sqrt ev else 0 + +-- ============================================================================ +-- Golub-Kahan bidiagonalisation SVD (GVL4 Algorithm 5.4.2 + 8.6.2) +-- ============================================================================ + +-- | Full Golub-Kahan SVD pipeline. +-- Phase 1: Bidiagonalise A → U₀ B V₀^T +-- Phase 2: Implicit-shift QR on bidiagonal B, accumulating rotations into U, V +-- Phase 3: Assemble final U, sigma, V; ensure σᵢ ≥ 0; sort descending +svdGKP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double + -> (Matrix m m M.P Double, Vector n M.P Double, Matrix n n M.P Double) +svdGKP (MkMatrix a_) = runST $ do + let !mm = dimVal @m + !nn = dimVal @n + + -- Copy input into mutable working storage + mA <- M.thawS a_ + let mbaA = unwrapMutableByteArray mA + offA = unwrapMutableByteArrayOffset mA + + -- Allocate arrays for Householder betas + mbaBetaL <- newByteArray (nn * 8) -- left Householder betas + mbaBetaR <- newByteArray (nn * 8) -- right Householder betas + + -- Phase 1: Bidiagonalise A in-place (BLAS-3 panel for large, Level-2 for small) + bidiagonalizePPanel mbaA offA mm nn mbaBetaL mbaBetaR + + -- Extract diagonal d and superdiagonal e from bidiagonalised A + mbaD <- newByteArray (nn * 8) + mbaE <- newByteArray (nn * 8) + forM_ [0..nn-1] $ \k -> do + dk <- readRawD mbaA offA (k * nn + k) + writeRawD mbaD 0 k dk + forM_ [0..nn-2] $ \k -> do + ek <- readRawD mbaA offA (k * nn + (k+1)) + writeRawD mbaE 0 k ek + + -- Freeze A for Householder vector extraction + frozenA <- M.freezeS mA + + -- Phase 2: Accumulate U₀ and V₀ from stored Householder vectors + -- U₀ = H₀ H₁ ... H_{n-1} (left reflectors, stored in columns of A) + mU <- M.newMArray @M.P (Sz (mm :. mm)) (0 :: Double) + let mbaU = unwrapMutableByteArray mU + offU = unwrapMutableByteArrayOffset mU + -- Initialise U = I + forM_ [0..mm-1] $ \i -> + writeRawD mbaU offU (i * mm + i) 1.0 + + let baA = unwrapByteArray frozenA + offFA = unwrapByteArrayOffset frozenA + + -- Accumulate left Householder reflectors into U (forward: U = H₀ H₁ ⋯ H_{n-1}) + -- Left reflector k: v stored in column k of A, rows k+1..m-1, with v[k]=1 implicit + if nn <= 16 + then + -- Small matrix: per-row accumulation (Level-2) + forM_ [0..nn-1] $ \k -> do + betaK <- readRawD mbaBetaL 0 k + when (betaK /= 0) $ + forM_ [0..mm-1] $ \row -> + rawMutQAccum mbaU offU mm baA offFA nn betaK k mm row + else do + -- Blocked WY: batch nb Householder vectors at a time + let !nbU = min 48 nn + mbaYU <- newByteArray (mm * nbU * 8) + mbaTfU <- newByteArray (nbU * nbU * 8) + mbaW1U <- newByteArray (mm * nbU * 8) + mbaW2U <- newByteArray (mm * nbU * 8) + mbaYTU <- newByteArray (nbU * mm * 8) + mbaGU <- newByteArray (nbU * nbU * 8) + + let goBlockU !k0 + | k0 >= nn = pure () + | otherwise = do + let !bsz = min nbU (nn - k0) + -- Pack Y (mm × bsz): Y[:,j] = left Householder vector k0+j + rawZeroDoubles mbaYU 0 (mm * bsz) + forM_ [0..bsz-1] $ \j -> do + let !k = k0 + j + writeRawD mbaYU 0 (k * bsz + j) 1.0 + forM_ [k+1..mm-1] $ \l -> + writeRawD mbaYU 0 (l * bsz + j) (readBA baA offFA (l * nn + k)) + + -- Transpose Y → Y^T (bsz × mm) for GEMM reuse + rawZeroDoubles mbaYTU 0 (bsz * mm) + forM_ [0..bsz-1] $ \j -> do + let !k = k0 + j + writeRawD mbaYTU 0 (j * mm + k) 1.0 + forM_ [k+1..mm-1] $ \l -> + writeRawD mbaYTU 0 (j * mm + l) (readBA baA offFA (l * nn + k)) + + baYU <- unsafeFreezeByteArray mbaYU + baYTU <- unsafeFreezeByteArray mbaYTU + + -- G = Y^T × Y (bsz × bsz) + rawZeroDoubles mbaGU 0 (bsz * bsz) + rawGemmKernel baYTU 0 baYU 0 mbaGU 0 bsz mm bsz + + -- Build T-factor (bsz × bsz upper-triangular) + rawZeroDoubles mbaTfU 0 (bsz * bsz) + forM_ [0..bsz-1] $ \j -> do + betaj <- readRawD mbaBetaL 0 (k0 + j) + writeRawD mbaTfU 0 (j * bsz + j) betaj + when (j > 0 && betaj /= 0) $ do + -- T[0..j-1, j] = -betaj * T[0..j-1, 0..j-1] * G[0..j-1, j] + forM_ [0..j-1] $ \i -> do + g_ij <- readRawD mbaGU 0 (i * bsz + j) + writeRawD mbaW1U 0 i g_ij + forM_ [0..j-1] $ \i -> do + let triLoop !l !acc + | l >= j = pure acc + | otherwise = do + til <- readRawD mbaTfU 0 (i * bsz + l) + dl <- readRawD mbaW1U 0 l + triLoop (l+1) (acc + til * dl) + z <- triLoop 0 0 + writeRawD mbaTfU 0 (i * bsz + j) (negate betaj * z) + + -- W1 = Q · Y (mm×mm * mm×bsz → mm×bsz) + baQU <- unsafeFreezeByteArray mbaU + rawZeroDoubles mbaW1U 0 (mm * bsz) + rawGemmKernel baQU offU baYU 0 mbaW1U 0 mm mm bsz + + -- W2 = W1 · T (mm×bsz * bsz×bsz → mm×bsz) + baW1U <- unsafeFreezeByteArray mbaW1U + baTfU <- unsafeFreezeByteArray mbaTfU + rawZeroDoubles mbaW2U 0 (mm * bsz) + rawGemmKernel baW1U 0 baTfU 0 mbaW2U 0 mm bsz bsz + + -- Negate W2 + rawNegateDoubles mbaW2U 0 (mm * bsz) + + -- Q += (-W2) · Y^T (mm×bsz * bsz×mm → mm×mm) + baNW2U <- unsafeFreezeByteArray mbaW2U + rawGemmKernel baNW2U 0 baYTU 0 mbaU offU mm bsz mm + + goBlockU (k0 + bsz) + goBlockU 0 + + -- V₀ = G₁ G₂ ... G_{n-3} (right reflectors) + mV <- M.newMArray @M.P (Sz (nn :. nn)) (0 :: Double) + let mbaV = unwrapMutableByteArray mV + offV = unwrapMutableByteArrayOffset mV + -- Initialise V = I + forM_ [0..nn-1] $ \i -> + writeRawD mbaV offV (i * nn + i) 1.0 + + -- Accumulate right Householder reflectors into V (forward: V = G₀ G₁ ⋯ G_{n-3}) + -- Right reflector k: v stored in row k of A, cols k+2..n-1, with v[k+1]=1 implicit + if nn < 19 + then + -- Small: per-row Level-2 + when (nn >= 3) $ + forM_ [0..nn-3] $ \k -> do + betaK <- readRawD mbaBetaR 0 k + when (betaK /= 0) $ + forM_ [0..nn-1] $ \row -> + rightQAccum mbaV offV nn baA offFA nn betaK k nn row + else when (nn >= 3) $ do + -- Blocked WY for right reflectors + let !nRefl = nn - 2 -- right reflectors 0..nn-3 + !nbV = min 48 nRefl + mbaYV <- newByteArray (nn * nbV * 8) + mbaTfV <- newByteArray (nbV * nbV * 8) + mbaW1V <- newByteArray (nn * nbV * 8) + mbaW2V <- newByteArray (nn * nbV * 8) + mbaYTV <- newByteArray (nbV * nn * 8) + mbaGV <- newByteArray (nbV * nbV * 8) + + let goBlockV !k0 + | k0 >= nRefl = pure () + | otherwise = do + let !bsz = min nbV (nRefl - k0) + -- Pack Y (nn × bsz): Y[:,j] = right Householder vector k0+j + -- Right vector k has implicit 1 at position k+1, stored values at k+2..nn-1 + rawZeroDoubles mbaYV 0 (nn * bsz) + forM_ [0..bsz-1] $ \j -> do + let !k = k0 + j + writeRawD mbaYV 0 ((k+1) * bsz + j) 1.0 + forM_ [k+2..nn-1] $ \l -> + writeRawD mbaYV 0 (l * bsz + j) (readBA baA offFA (k * nn + l)) + + -- Transpose Y → Y^T (bsz × nn) + rawZeroDoubles mbaYTV 0 (bsz * nn) + forM_ [0..bsz-1] $ \j -> do + let !k = k0 + j + writeRawD mbaYTV 0 (j * nn + (k+1)) 1.0 + forM_ [k+2..nn-1] $ \l -> + writeRawD mbaYTV 0 (j * nn + l) (readBA baA offFA (k * nn + l)) + + baYV <- unsafeFreezeByteArray mbaYV + baYTV <- unsafeFreezeByteArray mbaYTV + + -- G = Y^T × Y (bsz × bsz) + rawZeroDoubles mbaGV 0 (bsz * bsz) + rawGemmKernel baYTV 0 baYV 0 mbaGV 0 bsz nn bsz + + -- Build T-factor (bsz × bsz upper-triangular) + rawZeroDoubles mbaTfV 0 (bsz * bsz) + forM_ [0..bsz-1] $ \j -> do + betaj <- readRawD mbaBetaR 0 (k0 + j) + writeRawD mbaTfV 0 (j * bsz + j) betaj + when (j > 0 && betaj /= 0) $ do + -- T[0..j-1, j] = -betaj * T[0..j-1, 0..j-1] * G[0..j-1, j] + forM_ [0..j-1] $ \i -> do + g_ij <- readRawD mbaGV 0 (i * bsz + j) + writeRawD mbaW1V 0 i g_ij + forM_ [0..j-1] $ \i -> do + let triLoop !l !acc + | l >= j = pure acc + | otherwise = do + til <- readRawD mbaTfV 0 (i * bsz + l) + dl <- readRawD mbaW1V 0 l + triLoop (l+1) (acc + til * dl) + z <- triLoop 0 0 + writeRawD mbaTfV 0 (i * bsz + j) (negate betaj * z) + + -- W1 = V · Y (nn×nn * nn×bsz → nn×bsz) + baQV <- unsafeFreezeByteArray mbaV + rawZeroDoubles mbaW1V 0 (nn * bsz) + rawGemmKernel baQV offV baYV 0 mbaW1V 0 nn nn bsz + + -- W2 = W1 · T (nn×bsz * bsz×bsz → nn×bsz) + baW1V <- unsafeFreezeByteArray mbaW1V + baTfV <- unsafeFreezeByteArray mbaTfV + rawZeroDoubles mbaW2V 0 (nn * bsz) + rawGemmKernel baW1V 0 baTfV 0 mbaW2V 0 nn bsz bsz + + -- Negate W2 + rawNegateDoubles mbaW2V 0 (nn * bsz) + + -- V += (-W2) · Y^T (nn×bsz * bsz×nn → nn×nn) + baNW2V <- unsafeFreezeByteArray mbaW2V + rawGemmKernel baNW2V 0 baYTV 0 mbaV offV nn bsz nn + + goBlockV (k0 + bsz) + goBlockV 0 + + -- Phase 3: Bidiagonal SVD (D&C for large, QR iteration for small) + if nn >= dcBidiagThreshold + then dcBidiagSVD mbaD 0 mbaE 0 mbaU offU mm mbaV offV nn nn 1e-14 + else bidiagQRIterPCM mbaD 0 mbaE 0 mbaU offU mm mbaV offV nn nn (30 * nn) + + -- Phase 4: Ensure σᵢ ≥ 0 (flip sign of U column if needed) + forM_ [0..nn-1] $ \k -> do + dk <- readRawD mbaD 0 k + when (dk < 0) $ do + writeRawD mbaD 0 k (negate dk) + -- Flip column k of U + forM_ [0..mm-1] $ \i -> do + uik <- readRawD mbaU offU (i * mm + k) + writeRawD mbaU offU (i * mm + k) (negate uik) + + -- Phase 5: Sort singular values descending and permute U, V columns + pairs <- mapM (\k -> do dk <- readRawD mbaD 0 k; return (dk, k)) [0..nn-1] + let !sorted = sortBy (\(a1,_) (b1,_) -> compare (Down a1) (Down b1)) pairs + + frozenU <- M.freezeS mU + frozenV <- M.freezeS mV + let baU = unwrapByteArray frozenU + offFU = unwrapByteArrayOffset frozenU + baV = unwrapByteArray frozenV + offFV = unwrapByteArrayOffset frozenV + + let !sigmaVec = makeVector @n @M.P $ \i -> fst (sorted !! i) + !uMat = makeMatrix @m @m @M.P $ \i j -> + if j < nn + then let origCol = snd (sorted !! j) + in readBA baU offFU (i * mm + origCol) + else if i == j then 1 else 0 + !vMat = makeMatrix @n @n @M.P $ \i j -> + let origCol = snd (sorted !! j) + in readBA baV offFV (i * nn + origCol) + + return (uMat, sigmaVec, vMat) +{-# NOINLINE svdGKP #-} + +-- | In-place bidiagonalisation of an m×n matrix stored in a MutableByteArray. +-- GVL4 Algorithm 5.4.2, p. 284. +-- +-- After this, the matrix has: +-- - Diagonal d[k] = A[k,k] +-- - Superdiagonal e[k] = A[k,k+1] +-- - Left Householder vectors stored in column k below diagonal (rows k+1..m-1) +-- - Right Householder vectors stored in row k right of superdiag (cols k+2..n-1) +-- - Householder betas stored in mbaBetaL and mbaBetaR +bidiagonalizeP :: MutableByteArray s -> Int -> Int -> Int + -> MutableByteArray s -> MutableByteArray s -> ST s () +bidiagonalizeP mbaA offA mm nn mbaBetaL mbaBetaR = do + forM_ [0..nn-1] $ \k -> do + -- Left Householder: zero out A[k+1:m, k] + -- Compute Householder vector for column k, rows k..m-1 + if k < mm - 1 + then do + -- sigma = Σ A[i,k]² for i in k+1..m-1 + sigma <- rawMutSumSqColumn mbaA offA nn (k+1) mm k + x0 <- readRawD mbaA offA (k * nn + k) + if sigma < 1e-300 + then writeRawD mbaBetaL 0 k 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + -- Store v: normalise by v0 + -- v[k] will become 1 (implicit), v[k+1..m-1] = A[i,k]/v0 + forM_ [k+1..mm-1] $ \i -> do + aik <- readRawD mbaA offA (i * nn + k) + writeRawD mbaA offA (i * nn + k) (aik / v0) + -- Set A[k,k] = mu (the diagonal value after reflection) + writeRawD mbaA offA (k * nn + k) mu + writeRawD mbaBetaL 0 k beta + -- Apply left Householder to columns k+1..n-1 + -- Using rawMutHouseholderApply which reads v from column k, rows k+1..m-1 + forM_ [k+1..nn-1] $ \col -> + rawMutHouseholderApply mbaA offA nn beta k mm col + else + writeRawD mbaBetaL 0 k 0 + + -- Right Householder: zero out A[k, k+2:n] + if k < nn - 2 + then do + -- sigma = Σ A[k,j]² for j in k+2..n-1 + sigma <- rawMutSumSqRow mbaA offA nn k (k+2) nn + x0 <- readRawD mbaA offA (k * nn + (k+1)) + if sigma < 1e-300 + then writeRawD mbaBetaR 0 k 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + -- Store v: normalise by v0 + -- v[k+1] will become 1 (implicit), v[k+2..n-1] = A[k,j]/v0 + forM_ [k+2..nn-1] $ \j -> do + akj <- readRawD mbaA offA (k * nn + j) + writeRawD mbaA offA (k * nn + j) (akj / v0) + -- Set A[k,k+1] = mu (the superdiagonal value) + writeRawD mbaA offA (k * nn + (k+1)) mu + writeRawD mbaBetaR 0 k beta + -- Apply right Householder to rows k+1..m-1 + -- v is stored in row k, cols k+2..n-1, with implicit v[k+1]=1 + forM_ [k+1..mm-1] $ \row -> + rawMutHouseholderApplyRow mbaA offA nn beta k (k+1) nn row + else + when (k < nn - 1) $ writeRawD mbaBetaR 0 k 0 +{-# NOINLINE bidiagonalizeP #-} + +-- | BLAS-3 panel bidiagonalisation (DLABRD-style, GVL4 §5.4.3). +-- Processes nb columns at a time, deferring trailing updates via X, Y +-- accumulators and applying them as rank-nb GEMMs. +-- Falls back to Level-2 bidiagonalizeP for n < panelBidiagCrossover. +bidiagonalizePPanel :: MutableByteArray s -> Int -> Int -> Int + -> MutableByteArray s -> MutableByteArray s -> ST s () +bidiagonalizePPanel mbaA offA mm nn mbaBetaL mbaBetaR + | nn < panelBidiagCrossover = bidiagonalizeP mbaA offA mm nn mbaBetaL mbaBetaR + | otherwise = do + let !nb = min 32 (max 8 (nn `div` 6)) + -- Allocate accumulators: X (mm × nb), Y (nn × nb), row-major + mbaX <- newByteArray (mm * nb * 8) + mbaY <- newByteArray (nn * nb * 8) + -- Temp vectors for dot products + mbaZL <- newByteArray (nb * 8) -- V_L^T * v or Y^T * u + mbaZX <- newByteArray (nb * 8) -- X^T * v or V_R^T * u + -- Buffers for trailing GEMM + mbaVLbuf <- newByteArray (mm * nb * 8) + mbaYTbuf <- newByteArray (nb * nn * 8) + mbaXbuf <- newByteArray (mm * nb * 8) + mbaVRTbuf <- newByteArray (nb * nn * 8) + mbaTrail <- newByteArray (mm * nn * 8) + + let goPanel !k0 + | k0 >= nn - 1 = pure () + | otherwise = do + let !bs = min nb (nn - 1 - k0) + if bs < 2 -- last column: use Level-2 + then bidiagLastCols mbaA offA mm nn mbaBetaL mbaBetaR k0 + else do + rawZeroDoubles mbaX 0 (mm * bs) + rawZeroDoubles mbaY 0 (nn * bs) + -- Panel phase + panelBidiagStep mbaA offA mm nn mbaBetaL mbaBetaR + mbaX mbaY mbaZL mbaZX k0 bs + -- Trailing update + let !remR = mm - k0 - bs + !remC = nn - k0 - bs + when (remR > 0 && remC > 0) $ + applyTrailingUpdate mbaA offA mm nn mbaX mbaY + mbaVLbuf mbaYTbuf mbaXbuf mbaVRTbuf mbaTrail + k0 bs remR remC + goPanel (k0 + bs) + goPanel 0 + where + panelBidiagCrossover = 64 +{-# NOINLINE bidiagonalizePPanel #-} + +-- | Finish remaining columns with Level-2 bidiagonalisation. +bidiagLastCols :: MutableByteArray s -> Int -> Int -> Int + -> MutableByteArray s -> MutableByteArray s -> Int -> ST s () +bidiagLastCols mbaA offA mm nn mbaBetaL mbaBetaR k0 = do + forM_ [k0..nn-1] $ \k -> do + -- Left Householder + if k < mm - 1 + then do + sigma <- rawMutSumSqColumn mbaA offA nn (k+1) mm k + x0 <- readRawD mbaA offA (k * nn + k) + if sigma < 1e-300 + then writeRawD mbaBetaL 0 k 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + forM_ [k+1..mm-1] $ \i -> do + aik <- readRawD mbaA offA (i * nn + k) + writeRawD mbaA offA (i * nn + k) (aik / v0) + writeRawD mbaA offA (k * nn + k) mu + writeRawD mbaBetaL 0 k beta + forM_ [k+1..nn-1] $ \col -> + rawMutHouseholderApply mbaA offA nn beta k mm col + else writeRawD mbaBetaL 0 k 0 + -- Right Householder + if k < nn - 2 + then do + sigma <- rawMutSumSqRow mbaA offA nn k (k+2) nn + x0 <- readRawD mbaA offA (k * nn + (k+1)) + if sigma < 1e-300 + then writeRawD mbaBetaR 0 k 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + forM_ [k+2..nn-1] $ \j -> do + akj <- readRawD mbaA offA (k * nn + j) + writeRawD mbaA offA (k * nn + j) (akj / v0) + writeRawD mbaA offA (k * nn + (k+1)) mu + writeRawD mbaBetaR 0 k beta + forM_ [k+1..mm-1] $ \row -> + rawMutHouseholderApplyRow mbaA offA nn beta k (k+1) nn row + else when (k < nn - 1) $ writeRawD mbaBetaR 0 k 0 + +-- | DLABRD-style panel step: compute bs left/right Householder reflectors +-- starting at column k0, maintaining X and Y accumulators. +-- After this, A_eff[i,c] = A[i,c] - V_L[i,:]*Y[c,:]^T - X[i,:]*V_R[c,:]^T +-- for all i >= k0+bs, c >= k0+bs. +panelBidiagStep :: MutableByteArray s -> Int -> Int -> Int + -> MutableByteArray s -> MutableByteArray s + -> MutableByteArray s -> MutableByteArray s + -> MutableByteArray s -> MutableByteArray s + -> Int -> Int -> ST s () +panelBidiagStep mbaA offA mm nn mbaBetaL mbaBetaR mbaX mbaY mbaZL mbaZX k0 bs = do + forM_ [0..bs-1] $ \j -> do + let !k = k0 + j + + -- ================================================================ + -- PART A: Left Householder on column k + -- ================================================================ + + -- Step A1: Read corrected column k into A (in-place correction for rows k..m-1). + -- A_corr[i,k] = A[i,k] - sum_{l 0) $ + forM_ [k..mm-1] $ \i -> do + aik <- readRawD mbaA offA (i * nn + k) + -- V_L[i,l] * Y[k,l] sum + cVLY <- panelDot_VLY mbaA offA nn mbaY bs k0 i k j + -- X[i,l] * V_R[k,l] sum + cXVR <- panelDot_XVR mbaA offA nn mbaX bs k0 i k j + writeRawD mbaA offA (i * nn + k) (aik - cVLY - cXVR) + + -- Step A2: Left Householder from corrected column k, rows k..m-1 + if k < mm - 1 + then do + sigma <- rawMutSumSqColumn mbaA offA nn (k+1) mm k + x0 <- readRawD mbaA offA (k * nn + k) + if sigma < 1e-300 + then do + writeRawD mbaBetaL 0 k 0 + -- Zero Y column j + forM_ [0..nn-1] $ \c -> writeRawD mbaY 0 (c * bs + j) 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + -- Normalise and store HH vector in column k + forM_ [k+1..mm-1] $ \i -> do + aik <- readRawD mbaA offA (i * nn + k) + writeRawD mbaA offA (i * nn + k) (aik / v0) + writeRawD mbaA offA (k * nn + k) mu + writeRawD mbaBetaL 0 k beta + + -- Step A3: Compute Y column j + -- Precompute zL[l] = V_L[:,l]^T * v for l = 0..j-1 + forM_ [0..j-1] $ \l -> do + d <- dotVL_v mbaA offA nn k0 k mm l + writeRawD mbaZL 0 l d + -- Precompute zX[l] = X[:,l]^T * v for l = 0..j-1 + forM_ [0..j-1] $ \l -> do + d <- dotX_v mbaA offA nn mbaX bs k mm l + writeRawD mbaZX 0 l d + -- Y[c, j] = beta * (A^T*v[c] - sum_l Y[c,l]*zL[l] - sum_l V_R[c,l]*zX[l]) + forM_ [0..k] $ \c -> writeRawD mbaY 0 (c * bs + j) 0 + forM_ [k+1..nn-1] $ \c -> do + atv <- dotAT_v mbaA offA nn k mm c + ycorr <- dotAccum mbaY bs c mbaZL j + vrcorr <- dotVR_zX mbaA offA nn k0 mbaZX c j + writeRawD mbaY 0 (c * bs + j) (beta * (atv - ycorr - vrcorr)) + else do + writeRawD mbaBetaL 0 k 0 + forM_ [0..nn-1] $ \c -> writeRawD mbaY 0 (c * bs + j) 0 + + -- ================================================================ + -- PART B: Correct row k, then right Householder (if applicable) + -- ================================================================ + + -- Step B1: ALWAYS correct row k for columns k+1..n-1. + -- This is needed both for the right HH (if k < nn-2) and for the + -- superdiagonal entry e[k] = A[k, k+1] and trailing column values. + -- A_eff[k,c] = A[k,c] - sum_{l<=j} V_L[k,l]*Y[c,l] - sum_{l do + akc <- readRawD mbaA offA (k * nn + c) + cVLY <- panelDot_VLY mbaA offA nn mbaY bs k0 k c (j+1) + cXVR <- panelDot_XVR mbaA offA nn mbaX bs k0 k c j + writeRawD mbaA offA (k * nn + c) (akc - cVLY - cXVR) + + -- Step B2: Right Householder from corrected row k (only if k < nn-2) + if k < nn - 2 + then do + sigma <- rawMutSumSqRow mbaA offA nn k (k+2) nn + x0 <- readRawD mbaA offA (k * nn + (k+1)) + if sigma < 1e-300 + then do + writeRawD mbaBetaR 0 k 0 + forM_ [0..mm-1] $ \i -> writeRawD mbaX 0 (i * bs + j) 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + gamma = 2 * v0 * v0 / (sigma + v0 * v0) + forM_ [k+2..nn-1] $ \c -> do + akc <- readRawD mbaA offA (k * nn + c) + writeRawD mbaA offA (k * nn + c) (akc / v0) + writeRawD mbaA offA (k * nn + (k+1)) mu + writeRawD mbaBetaR 0 k gamma + + -- Step B3: Compute X column j + -- Precompute zL'[l] = Y[:,l]^T * u for l = 0..j + forM_ [0..j] $ \l -> do + d <- dotY_u mbaA offA nn mbaY bs k l + writeRawD mbaZL 0 l d + -- Precompute zX'[l] = V_R[:,l]^T * u for l = 0..j-1 + forM_ [0..j-1] $ \l -> do + d <- dotVR_u mbaA offA nn k0 k l + writeRawD mbaZX 0 l d + -- X[i, j] = gamma * (A*u[i] - sum_l V_L[i,l]*zL'[l] - sum_l X[i,l]*zX'[l]) + forM_ [0..k] $ \i -> writeRawD mbaX 0 (i * bs + j) 0 + forM_ [k+1..mm-1] $ \i -> do + au <- dotA_u mbaA offA nn k i + vlcorr <- dotVL_zL mbaA offA nn k0 mbaZL i (j+1) + xcorr <- dotX_zX mbaX bs mbaZX i j + writeRawD mbaX 0 (i * bs + j) (gamma * (au - vlcorr - xcorr)) + else do + when (k < nn - 1) $ writeRawD mbaBetaR 0 k 0 + forM_ [0..mm-1] $ \i -> writeRawD mbaX 0 (i * bs + j) 0 + +-- Helper: sum_l V_L[i,l]*Y[c,l] for l = 0..nL-1 +panelDot_VLY :: MutableByteArray s -> Int -> Int -> MutableByteArray s -> Int + -> Int -> Int -> Int -> Int -> ST s Double +panelDot_VLY mbaA offA nn mbaY bs k0 i c nL = go 0 0 + where + go !l !acc + | l >= nL = pure acc + | otherwise = do + let !kl = k0 + l + vl <- if i == kl then pure 1.0 + else if i > kl then readRawD mbaA offA (i * nn + kl) + else pure 0.0 + ycl <- readRawD mbaY 0 (c * bs + l) + go (l+1) (acc + vl * ycl) + +-- Helper: sum_l X[i,l]*V_R[c,l] for l = 0..nR-1 +panelDot_XVR :: MutableByteArray s -> Int -> Int -> MutableByteArray s -> Int + -> Int -> Int -> Int -> Int -> ST s Double +panelDot_XVR mbaA offA nn mbaX bs k0 i c nR = go 0 0 + where + go !l !acc + | l >= nR = pure acc + | otherwise = do + xil <- readRawD mbaX 0 (i * bs + l) + let !kl = k0 + l + vr <- if c == kl + 1 then pure 1.0 + else if c > kl + 1 then readRawD mbaA offA (kl * nn + c) + else pure 0.0 + go (l+1) (acc + xil * vr) + +-- Helper: V_L[:,l]^T * v where v = [1, A[k+1:m-1, k]] +dotVL_v :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> Int -> ST s Double +dotVL_v mbaA offA nn k0 k mm l = do + -- v[i-k]: v[0]=1, v[i-k]=A[i,k] for i>k + -- V_L[i,l]: 1 if i==kl, A[i,kl] if i>kl, 0 if i= k0+j and l < j, kl < k, so V_L[k,l] = A[k,kl] + vlk <- readRawD mbaA offA (k * nn + kl) + go (k+1) vlk + where + !kl = k0 + l + go !i !acc + | i >= mm = pure acc + | otherwise = do + vli <- readRawD mbaA offA (i * nn + kl) + vi <- readRawD mbaA offA (i * nn + k) + go (i+1) (acc + vli * vi) + +-- Helper: X[:,l]^T * v where v = [1, A[k+1:m-1, k]] +dotX_v :: MutableByteArray s -> Int -> Int -> MutableByteArray s -> Int + -> Int -> Int -> Int -> ST s Double +dotX_v mbaA offA nn mbaX bs k mm l = do + xkl <- readRawD mbaX 0 (k * bs + l) + go (k+1) xkl + where + go !i !acc + | i >= mm = pure acc + | otherwise = do + xil <- readRawD mbaX 0 (i * bs + l) + vi <- readRawD mbaA offA (i * nn + k) + go (i+1) (acc + xil * vi) + +-- Helper: A^T * v at column c, where v = [1, A[k+1:m-1, k]] +dotAT_v :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> ST s Double +dotAT_v mbaA offA nn k mm c = do + akc <- readRawD mbaA offA (k * nn + c) + go (k+1) akc + where + go !i !acc + | i >= mm = pure acc + | otherwise = do + aic <- readRawD mbaA offA (i * nn + c) + vi <- readRawD mbaA offA (i * nn + k) + go (i+1) (acc + aic * vi) + +-- Helper: sum_l Y[c,l]*zL[l] for l = 0..nL-1 +dotAccum :: MutableByteArray s -> Int -> Int -> MutableByteArray s -> Int -> ST s Double +dotAccum mbaY bs c mbaZL nL = go 0 0 + where + go !l !acc + | l >= nL = pure acc + | otherwise = do + ycl <- readRawD mbaY 0 (c * bs + l) + zl <- readRawD mbaZL 0 l + go (l+1) (acc + ycl * zl) + +-- Helper: sum_l V_R[c,l]*zX[l] for l = 0..nR-1 +dotVR_zX :: MutableByteArray s -> Int -> Int -> Int -> MutableByteArray s + -> Int -> Int -> ST s Double +dotVR_zX mbaA offA nn k0 mbaZX c nR = go 0 0 + where + go !l !acc + | l >= nR = pure acc + | otherwise = do + let !kl = k0 + l + vr <- if c == kl + 1 then pure 1.0 + else if c > kl + 1 then readRawD mbaA offA (kl * nn + c) + else pure 0.0 + zx <- readRawD mbaZX 0 l + go (l+1) (acc + vr * zx) + +-- Helper: Y[:,l]^T * u where u = [1, A[k, k+2:n-1]] +dotY_u :: MutableByteArray s -> Int -> Int -> MutableByteArray s -> Int + -> Int -> Int -> ST s Double +dotY_u mbaA offA nn mbaY bs k l = do + yk1l <- readRawD mbaY 0 ((k+1) * bs + l) + go (k+2) yk1l + where + go !c !acc + | c >= nn = pure acc + | otherwise = do + ycl <- readRawD mbaY 0 (c * bs + l) + uc <- readRawD mbaA offA (k * nn + c) + go (c+1) (acc + ycl * uc) + +-- Helper: V_R[:,l]^T * u where u = [1, A[k, k+2:n-1]] +dotVR_u :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> ST s Double +dotVR_u mbaA offA nn k0 k l = do + -- u[c-k-1]: u[0]=1 at c=k+1, u[c-k-1]=A[k,c] for c>k+1 + -- V_R[c,l]: 1 if c==kl+1, A[kl,c] if c>kl+1, 0 if c<=kl + -- We need sum_{c=k+1}^{n-1} V_R[c,l] * u[c-k-1] + -- Since k > kl (k=k0+j, l kl+1, i.e., k > kl) + vrkp1 <- if k + 1 == kl + 1 then pure 1.0 + else readRawD mbaA offA (kl * nn + (k+1)) + go (k+2) vrkp1 + where + !kl = k0 + l + go !c !acc + | c >= nn = pure acc + | otherwise = do + vrc <- if c == kl + 1 then pure 1.0 + else readRawD mbaA offA (kl * nn + c) + uc <- readRawD mbaA offA (k * nn + c) + go (c+1) (acc + vrc * uc) + +-- Helper: A * u at row i, where u = [1, A[k, k+2:n-1]] +dotA_u :: MutableByteArray s -> Int -> Int -> Int -> Int -> ST s Double +dotA_u mbaA offA nn k i = do + aikp1 <- readRawD mbaA offA (i * nn + (k+1)) + go (k+2) aikp1 + where + go !c !acc + | c >= nn = pure acc + | otherwise = do + aic <- readRawD mbaA offA (i * nn + c) + uc <- readRawD mbaA offA (k * nn + c) + go (c+1) (acc + aic * uc) + +-- Helper: sum_l V_L[i,l]*zL[l] for l = 0..nL-1 +dotVL_zL :: MutableByteArray s -> Int -> Int -> Int -> MutableByteArray s + -> Int -> Int -> ST s Double +dotVL_zL mbaA offA nn k0 mbaZL i nL = go 0 0 + where + go !l !acc + | l >= nL = pure acc + | otherwise = do + let !kl = k0 + l + vl <- if i == kl then pure 1.0 + else if i > kl then readRawD mbaA offA (i * nn + kl) + else pure 0.0 + zl <- readRawD mbaZL 0 l + go (l+1) (acc + vl * zl) + +-- Helper: sum_l X[i,l]*zX[l] for l = 0..nR-1 +dotX_zX :: MutableByteArray s -> Int -> MutableByteArray s -> Int -> Int -> ST s Double +dotX_zX mbaX bs mbaZX i nR = go 0 0 + where + go !l !acc + | l >= nR = pure acc + | otherwise = do + xil <- readRawD mbaX 0 (i * bs + l) + zx <- readRawD mbaZX 0 l + go (l+1) (acc + xil * zx) + +-- | Apply trailing GEMM update after a panel step. +-- A[k0+bs:m, k0+bs:n] -= V_L_trail * Y_trail^T + X_trail * V_R_trail^T +applyTrailingUpdate :: MutableByteArray s -> Int -> Int -> Int + -> MutableByteArray s -> MutableByteArray s + -> MutableByteArray s -> MutableByteArray s + -> MutableByteArray s -> MutableByteArray s + -> MutableByteArray s + -> Int -> Int -> Int -> Int -> ST s () +applyTrailingUpdate mbaA offA _mm nn mbaX mbaY + mbaVLbuf mbaYTbuf mbaXbuf mbaVRTbuf mbaTrail + k0 bs remR remC = do + let !trailRowStart = k0 + bs + !trailColStart = k0 + bs + + -- Copy A_trail to contiguous buffer (remR × remC) + forM_ [0..remR-1] $ \i -> + forM_ [0..remC-1] $ \c -> do + val <- readRawD mbaA offA ((trailRowStart + i) * nn + trailColStart + c) + writeRawD mbaTrail 0 (i * remC + c) val + + -- Build V_L_trail (remR × bs): V_L[trailRowStart+i, l] for i=0..remR-1, l=0..bs-1 + -- For all trail rows, i >= trailRowStart > k0+l, so V_L[i,l] = A[i, k0+l] + forM_ [0..remR-1] $ \i -> + forM_ [0..bs-1] $ \l -> + readRawD mbaA offA ((trailRowStart + i) * nn + (k0 + l)) >>= + writeRawD mbaVLbuf 0 (i * bs + l) + + -- Build Y_trail^T (bs × remC): Y_trail^T[l, c] = Y[trailColStart+c, l] + forM_ [0..bs-1] $ \l -> + forM_ [0..remC-1] $ \c -> + readRawD mbaY 0 ((trailColStart + c) * bs + l) >>= + writeRawD mbaYTbuf 0 (l * remC + c) + + -- GEMM 1: trail -= V_L_trail * Y_trail^T + -- Negate V_L_trail: nVL = -V_L_trail + rawNegateDoubles mbaVLbuf 0 (remR * bs) + baVL <- unsafeFreezeByteArray mbaVLbuf + baYT <- unsafeFreezeByteArray mbaYTbuf + rawGemmKernel baVL 0 baYT 0 mbaTrail 0 remR bs remC + + -- Build X_trail (remR × bs): X[trailRowStart+i, l] + forM_ [0..remR-1] $ \i -> + forM_ [0..bs-1] $ \l -> + readRawD mbaX 0 ((trailRowStart + i) * bs + l) >>= + writeRawD mbaXbuf 0 (i * bs + l) + + -- Build V_R_trail^T (bs × remC): V_R_trail^T[l, c] = V_R[trailColStart+c, l] + -- V_R[c, l] = 1 if c==k0+l+1, A[k0+l, c] if c>k0+l+1, 0 if c<=k0+l + -- Must handle implicit 1: when trailColStart+c == k0+l+1 (i.e., l=bs-1, c=0) + forM_ [0..bs-1] $ \l -> + forM_ [0..remC-1] $ \c -> do + let !globalC = trailColStart + c + !kl = k0 + l + val <- if globalC == kl + 1 then pure 1.0 + else if globalC > kl + 1 then readRawD mbaA offA (kl * nn + globalC) + else pure 0.0 + writeRawD mbaVRTbuf 0 (l * remC + c) val + + -- GEMM 2: trail -= X_trail * V_R_trail^T + rawNegateDoubles mbaXbuf 0 (remR * bs) + baX <- unsafeFreezeByteArray mbaXbuf + baVRT <- unsafeFreezeByteArray mbaVRTbuf + rawGemmKernel baX 0 baVRT 0 mbaTrail 0 remR bs remC + + -- Copy trail back to A + forM_ [0..remR-1] $ \i -> + forM_ [0..remC-1] $ \c -> do + val <- readRawD mbaTrail 0 (i * remC + c) + writeRawD mbaA offA ((trailRowStart + i) * nn + trailColStart + c) val + +-- | Implicit-shift bidiagonal QR iteration (GVL4 Algorithm 8.6.2). +-- Operates on diagonal d and superdiagonal e of an upper bidiagonal matrix. +-- Accumulates left rotations into U (m×n columns) and right rotations into V (n×n). +-- +-- Each iteration: (1) find the active unreduced block [p..q] by scanning from +-- the bottom for deflation, then scanning up for split; (2) apply one QR step +-- to [p..q]; (3) repeat until fully deflated or maxIter reached. +bidiagQRIterP :: MutableByteArray s -> Int -- d + offset + -> MutableByteArray s -> Int -- e + offset + -> MutableByteArray s -> Int -> Int -- U + offset + ucols + -> MutableByteArray s -> Int -> Int -- V + offset + vcols + -> Int -> Int -- n, maxIter + -> ST s () +bidiagQRIterP mbaD offD mbaE offE mbaU offU ucols mbaV offV vcols nn maxIter = go 0 + where + go !iter + | iter >= maxIter = return () + | otherwise = do + -- Step 1: Find q — the bottom of the unreduced block. + -- Scan from nn-1 downward, deflating negligible e[q-1]. + q <- deflateHi (nn - 1) + if q <= 0 + then return () -- fully deflated + else do + -- Step 2: Find p — the top of the unreduced block. + -- Scan from q-1 downward, looking for a split. + p <- findLo (q - 1) + -- Step 3: Apply one QR step to [p..q] + bidiagQRStep mbaD offD mbaE offE mbaU offU ucols mbaV offV vcols p q + go (iter + 1) + + -- Scan from hi down: deflate any trailing negligible superdiagonals. + -- Returns the index of the bottom row of the active block (0 if fully deflated). + deflateHi !hi + | hi <= 0 = return 0 + | otherwise = do + ehi <- readRawD mbaE offE (hi - 1) + dhi <- readRawD mbaD offD hi + dhi1 <- readRawD mbaD offD (hi - 1) + let tol = 1e-14 * (abs dhi1 + abs dhi) + if abs ehi <= tol + then do + writeRawD mbaE offE (hi - 1) 0 + deflateHi (hi - 1) + else return hi + + -- Scan from idx downward to find the top of the unreduced block. + -- Returns the smallest p such that B[p..q] is unreduced. + findLo !idx + | idx <= 0 = return 0 + | otherwise = do + eidx <- readRawD mbaE offE (idx - 1) + didx <- readRawD mbaD offD idx + didx1 <- readRawD mbaD offD (idx - 1) + let tol = 1e-14 * (abs didx1 + abs didx) + if abs eidx <= tol + then do + writeRawD mbaE offE (idx - 1) 0 + return idx + else findLo (idx - 1) +{-# NOINLINE bidiagQRIterP #-} + +-- | Column-major bidiagonal QR iteration with AED and stall detection. +-- Transposes U (mm×mm) and V (nn×nn) to column-major layout for SIMD Givens, +-- runs bidiag QR with aggressive early deflation, then transposes back. +-- Falls back to row-major path for nn < 10 (transpose overhead dominates). +bidiagQRIterPCM :: MutableByteArray s -> Int -- d + offset + -> MutableByteArray s -> Int -- e + offset + -> MutableByteArray s -> Int -> Int -- U + offset + ucols (= mm) + -> MutableByteArray s -> Int -> Int -- V + offset + vcols (= nn) + -> Int -> Int -- nn, maxIter + -> ST s () +bidiagQRIterPCM mbaD offD mbaE offE mbaU offU mm mbaV offV nn n maxIter + | n < 10 = bidiagQRIterP mbaD offD mbaE offE mbaU offU mm mbaV offV nn n maxIter + | otherwise = do + -- Transpose U (mm×mm) and V (nn×nn) to column-major + tmpU <- newByteArray (mm * mm * 8) + tmpV <- newByteArray (nn * nn * 8) + rawTransposeToColMajor mbaU offU tmpU 0 mm + rawTransposeToColMajor mbaV offV tmpV 0 nn + -- Run CM iteration + goCM 0 (n - 1) 0 + mbaD offD mbaE offE tmpU 0 mm tmpV 0 nn n maxIter + -- Transpose back to row-major + rawTransposeFromColMajor tmpU 0 mbaU offU mm + rawTransposeFromColMajor tmpV 0 mbaV offV nn +{-# NOINLINE bidiagQRIterPCM #-} + +-- | CM iteration core with AED and stall detection. +-- Parameters: iter, lastQ, stallCount, then the usual d/e/U/V arrays + n + maxIter. +goCM :: Int -> Int -> Int + -> MutableByteArray s -> Int -> MutableByteArray s -> Int + -> MutableByteArray s -> Int -> Int + -> MutableByteArray s -> Int -> Int + -> Int -> Int -> ST s () +goCM !iter !lastQ !stall mbaD offD mbaE offE mbaU offU mm mbaV offV nn n maxIter + | iter >= maxIter = return () + | stall >= 20 = return () -- stall detection: bail after 20 steps without deflation + | otherwise = do + -- AED: scan bottom w superdiagonal entries for aggressive early deflation + let w = min 6 ((n + 2) `div` 3) + aedScan (n - 1) w + -- Find q — bottom of unreduced block + q <- defHiCM (n - 1) + if q <= 0 + then return () -- fully deflated + else do + -- Find p — top of unreduced block + p <- findLoCM (q - 1) + -- Apply one CM QR step to [p..q] + bidiagQRStepCM mbaD offD mbaE offE mbaU offU mm mbaV offV nn p q + let !newStall = if q == lastQ then stall + 1 else 0 + goCM (iter + 1) q newStall mbaD offD mbaE offE mbaU offU mm mbaV offV nn n maxIter + where + -- AED: scan bottom w entries, deflating negligible superdiagonals + aedScan _ 0 = return () + aedScan k remaining + | k <= 0 = return () + | otherwise = do + ek <- readRawD mbaE offE (k - 1) + dk <- readRawD mbaD offD k + dk1 <- readRawD mbaD offD (k - 1) + let tol = 1e-14 * (abs dk1 + abs dk) + if abs ek <= tol + then do + writeRawD mbaE offE (k - 1) 0 + aedScan (k - 1) (remaining - 1) + else return () -- stop at first non-negligible entry + + defHiCM !hi + | hi <= 0 = return 0 + | otherwise = do + ehi <- readRawD mbaE offE (hi - 1) + dhi <- readRawD mbaD offD hi + dhi1 <- readRawD mbaD offD (hi - 1) + let tol = 1e-14 * (abs dhi1 + abs dhi) + if abs ehi <= tol + then do + writeRawD mbaE offE (hi - 1) 0 + defHiCM (hi - 1) + else return hi + + findLoCM !idx + | idx <= 0 = return 0 + | otherwise = do + eidx <- readRawD mbaE offE (idx - 1) + didx <- readRawD mbaD offD idx + didx1 <- readRawD mbaD offD (idx - 1) + let tol = 1e-14 * (abs didx1 + abs didx) + if abs eidx <= tol + then do + writeRawD mbaE offE (idx - 1) 0 + return idx + else findLoCM (idx - 1) +{-# NOINLINE goCM #-} + +-- | One implicit-shift QR step on bidiagonal [lo..hi] using column-major U,V. +-- Same as bidiagQRStep but calls rawMutApplyGivensColumnsCM for SIMD. +bidiagQRStepCM :: MutableByteArray s -> Int -- d + offset + -> MutableByteArray s -> Int -- e + offset + -> MutableByteArray s -> Int -> Int -- U_CM + offset + mm + -> MutableByteArray s -> Int -> Int -- V_CM + offset + nn + -> Int -> Int -- lo, hi + -> ST s () +bidiagQRStepCM mbaD offD mbaE offE mbaU offU mm mbaV offV nn lo hi = do + -- Compute Wilkinson shift from trailing 2×2 of T = B^T B + dhi1 <- readRawD mbaD offD (hi - 1) + dhi <- readRawD mbaD offD hi + ehi1 <- readRawD mbaE offE (hi - 1) + ehi2 <- if hi >= 2 then readRawD mbaE offE (hi - 2) else return 0 + + let t11 = dhi1 * dhi1 + (if hi - 1 > lo then ehi2 * ehi2 else 0) + t12 = dhi1 * ehi1 + t22 = dhi * dhi + ehi1 * ehi1 + delta = (t11 - t22) / 2 + signD = if delta >= 0 then 1 else -1 + mu = t22 - t12 * t12 / (delta + signD * sqrt (delta * delta + t12 * t12)) + + dlo <- readRawD mbaD offD lo + elo <- readRawD mbaE offE lo + let y = dlo * dlo - mu + z = dlo * elo + + goChase lo y z + where + goChase k y_ z_ = do + let (cosR, sinR) = givens y_ z_ + dk <- readRawD mbaD offD k + ek <- readRawD mbaE offE k + dk1 <- readRawD mbaD offD (k + 1) + + let dk' = cosR * dk + sinR * ek + ek' = -sinR * dk + cosR * ek + bulgeL = sinR * dk1 + dk1' = cosR * dk1 + + writeRawD mbaD offD k dk' + writeRawD mbaE offE k ek' + writeRawD mbaD offD (k + 1) dk1' + + -- Update e[k-1]: right Givens rotates entry from row above + when (k > lo) $ + writeRawD mbaE offE (k - 1) (cosR * y_ + sinR * z_) + + -- Accumulate right rotation into V (column-major, SIMD) + rawMutApplyGivensColumnsCM mbaV offV nn cosR sinR k (k+1) nn + + let (cosL, sinL) = givens dk' bulgeL + + let dk'' = cosL * dk' + sinL * bulgeL + ek'' = cosL * ek' + sinL * dk1' + dk1'' = -sinL * ek' + cosL * dk1' + + writeRawD mbaD offD k dk'' + writeRawD mbaE offE k ek'' + writeRawD mbaD offD (k + 1) dk1'' + + when (k + 1 < hi) $ do + ek1 <- readRawD mbaE offE (k + 1) + let bulgeR = sinL * ek1 + ek1' = cosL * ek1 + writeRawD mbaE offE (k + 1) ek1' + + -- Accumulate left rotation into U (column-major, SIMD) + rawMutApplyGivensColumnsCM mbaU offU mm cosL sinL k (k+1) mm + + goChase (k + 1) ek'' bulgeR + + when (k + 1 >= hi) $ + rawMutApplyGivensColumnsCM mbaU offU mm cosL sinL k (k+1) mm + +-- | One implicit-shift QR step on bidiagonal [lo..hi]. +-- Computes Wilkinson shift from bottom 2×2 of B^T B, +-- then chases bulge via Givens rotations. +bidiagQRStep :: MutableByteArray s -> Int -- d + offset + -> MutableByteArray s -> Int -- e + offset + -> MutableByteArray s -> Int -> Int -- U + offset + ucols + -> MutableByteArray s -> Int -> Int -- V + offset + vcols + -> Int -> Int -- lo, hi + -> ST s () +bidiagQRStep mbaD offD mbaE offE mbaU offU ucols mbaV offV vcols lo hi = do + -- Compute Wilkinson shift from trailing 2×2 of T = B^T B + dhi1 <- readRawD mbaD offD (hi - 1) + dhi <- readRawD mbaD offD hi + ehi1 <- readRawD mbaE offE (hi - 1) + ehi2 <- if hi >= 2 then readRawD mbaE offE (hi - 2) else return 0 + + -- T = B^T B trailing 2×2: + -- t11 = d[hi-1]^2 + e[hi-2]^2 (e[hi-2] = 0 if hi-1 == lo) + -- t12 = d[hi-1] * e[hi-1] + -- t22 = d[hi]^2 + e[hi-1]^2 + let t11 = dhi1 * dhi1 + (if hi - 1 > lo then ehi2 * ehi2 else 0) + t12 = dhi1 * ehi1 + t22 = dhi * dhi + ehi1 * ehi1 + + -- Wilkinson shift: eigenvalue of [[t11,t12],[t12,t22]] closer to t22 + let delta = (t11 - t22) / 2 + signD = if delta >= 0 then 1 else -1 + mu = t22 - t12 * t12 / (delta + signD * sqrt (delta * delta + t12 * t12)) + + -- Initial values for bulge chase + dlo <- readRawD mbaD offD lo + elo <- readRawD mbaE offE lo + let y = dlo * dlo - mu + z = dlo * elo + + -- Chase bulge from lo to hi + go lo y z + where + go k y_ z_ = do + -- Right Givens rotation G(k,k+1,θ) to zero z in [y; z] + let (cosR, sinR) = givens y_ z_ + -- Apply to columns k, k+1 of B (affects d[k], e[k], d[k+1], and possibly e[k-1]) + dk <- readRawD mbaD offD k + ek <- readRawD mbaE offE k + dk1 <- readRawD mbaD offD (k + 1) + + -- B * G^T: columns k and k+1 get mixed + let dk' = cosR * dk + sinR * ek + ek' = -sinR * dk + cosR * ek + -- This creates a bulge at B[k+1,k] + bulgeL = sinR * dk1 + dk1' = cosR * dk1 + + writeRawD mbaD offD k dk' + writeRawD mbaE offE k ek' + writeRawD mbaD offD (k + 1) dk1' + + -- Update e[k-1]: the right Givens also rotates the entry from the row above. + -- For k > lo: B[k-1,k] was y_, B[k-1,k+1] was z_ (the bulge). + -- After rotation: B[k-1,k] = cosR*y_ + sinR*z_ (= r), B[k-1,k+1] = 0. + when (k > lo) $ + writeRawD mbaE offE (k - 1) (cosR * y_ + sinR * z_) + + -- Accumulate right rotation into V (columns k, k+1) + rawMutApplyGivensColumns mbaV offV vcols cosR sinR k (k+1) vcols + + -- Left Givens rotation G(k,k+1,θ) to zero the bulge at (k+1, k) + let (cosL, sinL) = givens dk' bulgeL + + -- G * B: rows k and k+1 get mixed + let dk'' = cosL * dk' + sinL * bulgeL + ek'' = cosL * ek' + sinL * dk1' + dk1'' = -sinL * ek' + cosL * dk1' + + writeRawD mbaD offD k dk'' + writeRawD mbaE offE k ek'' + writeRawD mbaD offD (k + 1) dk1'' + + -- This may create a new bulge at position (k, k+2) if k+1 < hi + when (k + 1 < hi) $ do + ek1 <- readRawD mbaE offE (k + 1) + let bulgeR = sinL * ek1 + ek1' = cosL * ek1 + writeRawD mbaE offE (k + 1) ek1' + + -- Accumulate left rotation into U (columns k, k+1) + rawMutApplyGivensColumns mbaU offU ucols cosL sinL k (k+1) ucols + + -- Continue chase + go (k + 1) ek'' bulgeR + + when (k + 1 >= hi) $ + -- Accumulate final left rotation + rawMutApplyGivensColumns mbaU offU ucols cosL sinL k (k+1) ucols + +-- | Compute Givens rotation coefficients (c, s) such that +-- @[c, s; -s, c] [a; b] = [r; 0]@, i.e. @-s*a + c*b = 0@ and @r = c*a + s*b > 0@. +-- +-- This convention is chosen so that the bidiag QR bulge-chase formulas +-- @dk' = c*dk + s*ek@, @ek' = -s*dk + c*ek@ etc. are directly correct +-- for both left and right Givens rotations (GVL4 Algorithm 8.6.2). +givens :: Double -> Double -> (Double, Double) +givens a b + | b == 0 = (1, 0) + | abs b > abs a = + let tau = a / b + s = 1 / sqrt (1 + tau * tau) + c = s * tau + in (c, s) + | otherwise = + let tau = b / a + c = 1 / sqrt (1 + tau * tau) + s = c * tau + in (c, s) +{-# INLINE givens #-} + +-- | Accumulate a right Householder reflector into V. +-- Right reflector k: v stored in row k of frozen A, cols k+2..n-1, with v[k+1]=1 (implicit). +-- V = V * (I - beta * v * v^T) +-- For each row of V: V[row, k+1..n-1] -= (beta * Σ V[row,l] * v[l]) * v +rightQAccum :: MutableByteArray s -> Int -> Int -- V + offset + vcols + -> ByteArray -> Int -> Int -- frozen A + offset + acols + -> Double -> Int -> Int -> Int -- beta, k, n, row + -> ST s () +rightQAccum mbaV offV vcols (ByteArray baA) offFA acols beta k nn row = do + -- Phase 1: wi = beta * (V[row, k+1] + Σ_{l=k+2}^{n-1} V[row, l] * A[k, l]) + qrk1 <- readRawD mbaV offV (row * vcols + (k+1)) + acc <- goSum (k+2) 0 + let wi = beta * (qrk1 + acc) + -- Phase 2: V[row, k+1] -= wi (implicit v[k+1]=1) + writeRawD mbaV offV (row * vcols + (k+1)) (qrk1 - wi) + -- V[row, l] -= wi * A[k, l] for l in k+2..n-1 + goUpdate (k+2) wi + where + goSum l acc_ + | l >= nn = return acc_ + | otherwise = do + let vl = readBAI baA offFA (k * acols + l) + qrl <- readRawD mbaV offV (row * vcols + l) + goSum (l + 1) (acc_ + qrl * vl) + + goUpdate l wi + | l >= nn = return () + | otherwise = do + let vl = readBAI baA offFA (k * acols + l) + qrl <- readRawD mbaV offV (row * vcols + l) + writeRawD mbaV offV (row * vcols + l) (qrl - wi * vl) + goUpdate (l + 1) wi + + readBAI ba_ off_ i_ = + case indexDoubleArray# ba_ (case off_ of I# o -> o +# case i_ of I# ii -> ii) of + v -> D# v +{-# INLINE rightQAccum #-} + +-- | Read a Double from a MutableByteArray at element index. +readRawD :: MutableByteArray s -> Int -> Int -> ST s Double +readRawD (MutableByteArray mba) (I# off) (I# i) = ST $ \s -> + case readDoubleArray# mba (off +# i) s of (# s', v #) -> (# s', D# v #) +{-# INLINE readRawD #-} + +-- | Write a Double to a MutableByteArray at element index. +writeRawD :: MutableByteArray s -> Int -> Int -> Double -> ST s () +writeRawD (MutableByteArray mba) (I# off) (I# i) (D# v) = ST $ \s -> + case writeDoubleArray# mba (off +# i) v s of s' -> (# s', () #) +{-# INLINE writeRawD #-} + +-- ============================================================================ +-- Divide-and-conquer bidiagonal SVD (Gu-Eisenstat 1995, cf. LAPACK DBDSDC) +-- ============================================================================ + +-- | Small-subproblem threshold for D&C bidiagonal SVD. +-- Subproblems at or below this size use bidiag QR iteration. +dcBidiagThreshold :: Int +dcBidiagThreshold = 32 + +-- | Direct 2×2 upper bidiagonal SVD. +-- Given [[d0, e0], [0, d1]], compute singular values and rotation angles. +-- Returns (sigma_large, sigma_small, c_left, s_left, c_right, s_right) where +-- the left and right Givens rotations diagonalise B. +bidiag2x2SVD :: Double -> Double -> Double -> (Double, Double, Double, Double, Double, Double) +bidiag2x2SVD d0 e0 d1 = + -- B^T B = [[d0², d0*e0], [d0*e0, e0²+d1²]] + -- Eigenvalues of this 2×2 symmetric matrix give σ² + let !a11 = d0 * d0 + !a12 = d0 * e0 + !a22 = e0 * e0 + d1 * d1 + !tr = a11 + a22 + !det = a11 * a22 - a12 * a12 + !disc = max 0 (tr * tr - 4 * det) + !sqrtDisc = sqrt disc + !lam1 = (tr + sqrtDisc) / 2 + !lam2 = (tr - sqrtDisc) / 2 + !sig1 = sqrt (max 0 lam1) + !sig2 = sqrt (max 0 lam2) + -- Right rotation angle: diagonalise B^T B + -- (a11 - lam2) * cv + a12 * sv = 0 => cv/sv = -a12/(a11-lam2) + -- Or equivalently: atan2(a12, lam1 - a22) + !cv = if abs a12 < 1e-300 + then 1 + else let d = a11 - lam2 + r = sqrt (d * d + a12 * a12) + in d / r + !sv = if abs a12 < 1e-300 + then 0 + else let d = a11 - lam2 + r = sqrt (d * d + a12 * a12) + in a12 / r + -- Left rotation: B * V = U * Sigma + -- (d0*cv + e0*sv, -d0*sv + e0*cv) = (sig1*cu, sig2*su_neg) + -- (d1*sv, d1*cv) = (sig1*(-su), sig2*cu) + -- From first row: sig1*cu = d0*cv + e0*sv + !bv00 = d0 * cv + e0 * sv + !bv10 = d1 * sv + !r_l = sqrt (bv00 * bv00 + bv10 * bv10) + !cu = if r_l < 1e-300 then 1 else bv00 / r_l + !su = if r_l < 1e-300 then 0 else bv10 / r_l + in (sig1, sig2, cu, su, cv, sv) +{-# INLINE bidiag2x2SVD #-} + +-- | Divide-and-conquer bidiagonal SVD. +-- Replaces bidiagQRIterPCM for computing the SVD of a bidiagonal matrix. +-- +-- Input: d[0..nn-1] (diagonal), e[0..nn-2] (superdiagonal). +-- Output: d[] overwritten with singular values, +-- U (mm×mm) and V (nn×nn) updated with accumulated rotations. +dcBidiagSVD :: forall s. MutableByteArray s -> Int -- d + offset + -> MutableByteArray s -> Int -- e + offset + -> MutableByteArray s -> Int -> Int -- U + offset + mm + -> MutableByteArray s -> Int -> Int -- V + offset + nn + -> Int -- nn (bidiag dimension) + -> Double -- tolerance + -> ST s () +dcBidiagSVD mbaD offD mbaE offE mbaU offU mm mbaV offV nn0 nn tol = do + -- Pre-allocate all workspace once at maximum size + let !maxN = nn + wsLam <- newByteArray (maxN * 8) -- new eigenvalues (squared) + wsZ <- newByteArray (maxN * 8) -- z-vector + wsDSort <- newByteArray (maxN * 8) -- sorted d² values + wsZSort <- newByteArray (maxN * 8) -- sorted z values + wsDOrig <- newByteArray (maxN * 8) -- original d values (unsquared) + wsIdx <- newByteArray (maxN * 8) -- sort permutation (stored as Double) + wsPerm <- newByteArray (maxN * 8) -- deflation permutation (Int) + wsW <- newByteArray (maxN * maxN * 8) -- V-eigenvector matrix W_V + wsWU <- newByteArray (maxN * maxN * 8) -- U-eigenvector matrix W_U + + -- Local accumulators for V (nn×nn) and U (nn×nn) + -- U-local is nn×nn because we track rotations in singular-value index space + wsVlocal <- newByteArray (maxN * maxN * 8) + wsUlocal <- newByteArray (maxN * maxN * 8) + + -- GEMM workspace + wsVsub <- newByteArray (maxN * maxN * 8) -- V column extraction buffer + wsVres <- newByteArray (maxN * maxN * 8) -- V GEMM result + wsUsub <- newByteArray (maxN * maxN * 8) -- U column extraction buffer + wsUres <- newByteArray (maxN * maxN * 8) -- U GEMM result + wsQtemp <- newByteArray (maxN * maxN * 8) -- QR base case scratch + + -- Initialise local accumulators as identity + rawZeroDoubles wsVlocal 0 (maxN * maxN) + rawZeroDoubles wsUlocal 0 (maxN * maxN) + forM_ [0..maxN-1] $ \i -> do + writeRawD wsVlocal 0 (i * maxN + i) 1 + writeRawD wsUlocal 0 (i * maxN + i) 1 + + let -- Convert global index to local + toLocal g = g + + -- Apply a k×k rotation matrix to wsVlocal columns [colOff..colOff+k-1] + applyRotToVlocal !colOff !k rotMat = do + forM_ [0..k-1] $ \j -> + rawCopyColumn wsVlocal 0 maxN (colOff + j) wsVsub 0 k j maxN + baVsub <- unsafeFreezeByteArray wsVsub + baRot <- unsafeFreezeByteArray rotMat + rawZeroDoubles wsVres 0 (maxN * k) + rawGemmKernel baVsub 0 baRot 0 wsVres 0 maxN k k + forM_ [0..k-1] $ \j -> + rawCopyColumn wsVres 0 k j wsVlocal 0 maxN (colOff + j) maxN + + -- Apply a k×k rotation matrix to wsUlocal columns [colOff..colOff+k-1] + applyRotToUlocal !colOff !k rotMat = do + forM_ [0..k-1] $ \j -> + rawCopyColumn wsUlocal 0 maxN (colOff + j) wsUsub 0 k j maxN + baUsub <- unsafeFreezeByteArray wsUsub + baRot <- unsafeFreezeByteArray rotMat + rawZeroDoubles wsUres 0 (maxN * k) + rawGemmKernel baUsub 0 baRot 0 wsUres 0 maxN k k + forM_ [0..k-1] $ \j -> + rawCopyColumn wsUres 0 k j wsUlocal 0 maxN (colOff + j) maxN + + -- The recursive D&C function + dcGo :: Int -> Int -> ST s () -- s from ScopedTypeVariables + dcGo lo hi + -- Trivial: single element + | lo >= hi = return () + + -- Base case: 2×2 direct SVD + | hi == lo + 1 = do + d0_ <- readRawD mbaD offD lo + e0_ <- readRawD mbaE offE lo + d1_ <- readRawD mbaD offD hi + let (!sig1, !sig2, !cu, !su, !cv, !sv) = bidiag2x2SVD d0_ e0_ d1_ + writeRawD mbaD offD lo sig1 + writeRawD mbaD offD hi sig2 + writeRawD mbaE offE lo 0 + -- Apply left Givens to Ulocal columns + let !loL = toLocal lo + !hiL = toLocal hi + rawMutApplyGivensColumns wsUlocal 0 maxN cu su loL hiL maxN + -- Apply right Givens to Vlocal columns + rawMutApplyGivensColumns wsVlocal 0 maxN cv sv loL hiL maxN + + -- Small subproblem: use bidiag QR + GEMM to local accumulators + | hi - lo + 1 <= dcBidiagThreshold = do + let !k = hi - lo + 1 + !loL = toLocal lo + -- Initialise k×k identities for U and V rotations + rawZeroDoubles wsQtemp 0 (k * k) + rawZeroDoubles wsW 0 (k * k) + forM_ [0..k-1] $ \i -> do + writeRawD wsQtemp 0 (i * k + i) 1 + writeRawD wsW 0 (i * k + i) 1 + -- Run bidiag QR iteration: wsQtemp accumulates left, wsW accumulates right + bidiagQRIterP mbaD (offD + lo) mbaE (offE + lo) wsQtemp 0 k wsW 0 k k (30 * k) + -- Apply rotations to local accumulators via GEMM + applyRotToUlocal loL k wsQtemp + applyRotToVlocal loL k wsW + + -- D&C merge + | otherwise = do + let !k = (lo + hi) `div` 2 + !n1 = k - lo + 1 + !n2 = hi - k + !nn_ = hi - lo + 1 + !kL = toLocal k + !loL = toLocal lo + + -- Read and modify the coupling element + beta <- readRawD mbaE offE k + dk <- readRawD mbaD offD k + dk1 <- readRawD mbaD offD (k + 1) + let !absBeta = abs beta + !rho = absBeta + writeRawD mbaD offD k (dk - absBeta) + writeRawD mbaD offD (k+1) (dk1 - absBeta) + writeRawD mbaE offE k 0 + + -- Recurse on left [lo..k] and right [k+1..hi] subproblems + dcGo lo k + dcGo (k + 1) hi + + -- === Merge phase === + + -- Extract z-vector from Vlocal accumulator rows + -- z[0..n1-1] = last row of V₁ = row kL, columns loL..loL+n1-1 + forM_ [0..n1-1] $ \i -> do + qv <- readRawD wsVlocal 0 (kL * maxN + (loL + i)) + writeRawD wsZ 0 i qv + -- z[n1..nn_-1] = first row of V₂ = row (kL+1), columns loL+n1..loL+nn_-1 + forM_ [0..n2-1] $ \i -> do + qv <- readRawD wsVlocal 0 ((kL + 1) * maxN + (loL + n1 + i)) + let !zv = if beta < 0 then negate qv else qv + writeRawD wsZ 0 (n1 + i) zv + + -- Save original d-values (unsquared) for U-eigenvector computation + forM_ [0..nn_-1] $ \i -> do + di <- readRawD mbaD offD (lo + i) + writeRawD wsDOrig 0 i di + + -- Square d-values for secular equation and copy into sort buffers + forM_ [0..nn_-1] $ \i -> do + di <- readRawD mbaD offD (lo + i) + writeRawD wsDSort 0 i (di * di) + writeRawD wsZSort 0 i =<< readRawD wsZ 0 i + writeRawD wsIdx 0 i (fromIntegral i) + + -- Sort by d² values (insertion sort) + forM_ [1..nn_-1] $ \i -> do + di <- readRawD wsDSort 0 i + zi <- readRawD wsZSort 0 i + idxi <- readRawD wsIdx 0 i + dOi <- readRawD wsDOrig 0 i + let insertAt !j + | j < 0 = do + writeRawD wsDSort 0 0 di + writeRawD wsZSort 0 0 zi + writeRawD wsIdx 0 0 idxi + writeRawD wsDOrig 0 0 dOi + | otherwise = do + dj <- readRawD wsDSort 0 j + if dj > di + then do + writeRawD wsDSort 0 (j+1) dj + writeRawD wsZSort 0 (j+1) =<< readRawD wsZSort 0 j + writeRawD wsIdx 0 (j+1) =<< readRawD wsIdx 0 j + writeRawD wsDOrig 0 (j+1) =<< readRawD wsDOrig 0 j + insertAt (j - 1) + else do + writeRawD wsDSort 0 (j+1) di + writeRawD wsZSort 0 (j+1) zi + writeRawD wsIdx 0 (j+1) idxi + writeRawD wsDOrig 0 (j+1) dOi + insertAt (i - 1) + + -- Close-d deflation on squared values (cf. LAPACK dlaed2) + dMaxSq <- readRawD wsDSort 0 (nn_ - 1) + dMinSq <- readRawD wsDSort 0 0 + let !closeDTol = 8 * 2.220446049250313e-16 + * max (abs dMaxSq) (abs dMinSq + rho) + forM_ [0..nn_-2] $ \i -> do + di <- readRawD wsDSort 0 i + di1 <- readRawD wsDSort 0 (i + 1) + when (abs (di1 - di) <= closeDTol) $ do + zi <- readRawD wsZSort 0 i + zi1 <- readRawD wsZSort 0 (i + 1) + let !r = sqrt (zi * zi + zi1 * zi1) + when (r > 1e-300) $ do + let !c = zi1 / r + !s = zi / r + writeRawD wsZSort 0 i 0 + writeRawD wsZSort 0 (i + 1) r + -- Apply Givens to Vlocal columns (same as tridiagonal) + origI <- readRawD wsIdx 0 i + origI1 <- readRawD wsIdx 0 (i + 1) + let !colI = loL + (round origI :: Int) + !colI1 = loL + (round origI1 :: Int) + rawMutApplyGivensColumns wsVlocal 0 maxN c s colI colI1 maxN + -- Also apply to Ulocal + rawMutApplyGivensColumns wsUlocal 0 maxN c s colI colI1 maxN + + -- Perturbation-based deflation + zn2 <- sumZSq wsZSort 0 nn_ + let !eps_ = 2.220446049250313e-16 + !matNorm = max (abs dMaxSq) (abs dMinSq) + rho * zn2 + !basicDeflTol = max (tol * sqrt zn2) (8 * eps_ * matNorm) + !pertDeflTol = sqrt (eps_ * (1 + matNorm) / max rho 1e-300) + !deflTol = max basicDeflTol pertDeflTol + kND <- deflatePartition wsZSort 0 wsPerm 0 nn_ deflTol + + -- Extract Vlocal columns permuted by sort order into wsVsub (maxN × nn_) + forM_ [0..nn_-1] $ \sortedJ -> do + origIdx <- readRawD wsIdx 0 sortedJ + let !origJ = round origIdx :: Int + !srcCol = loL + origJ + rawCopyColumn wsVlocal 0 maxN srcCol wsVsub 0 nn_ sortedJ maxN + + -- Also extract Ulocal columns + forM_ [0..nn_-1] $ \sortedJ -> do + origIdx <- readRawD wsIdx 0 sortedJ + let !origJ = round origIdx :: Int + !srcCol = loL + origJ + rawCopyColumn wsUlocal 0 maxN srcCol wsUsub 0 nn_ sortedJ maxN + + if kND == 0 + then do + -- All deflated: singular values from sorted d², vectors from sorted cols + forM_ [0..nn_-1] $ \i -> do + rawCopyColumn wsVsub 0 nn_ i wsVlocal 0 maxN (loL + i) maxN + rawCopyColumn wsUsub 0 nn_ i wsUlocal 0 maxN (loL + i) maxN + forM_ [0..nn_-1] $ \i -> do + dsq <- readRawD wsDSort 0 i + writeRawD mbaD offD (lo + i) (sqrt (max 0 dsq)) + + else if kND == nn_ + then do + -- No deflation: full secular solve + eigenvectors + dual GEMM + secularSolve wsLam 0 wsDSort 0 wsZSort 0 rho nn_ deflTol + + -- V-eigenvectors via Gu-Eisenstat on (d², z, μ) + dcEigenvectors wsW 0 wsDSort 0 wsZSort 0 wsLam 0 rho nn_ + + -- U-eigenvectors: W_U[j,i] = dOrig[j] * W_V[j,i], then normalize + dcEigenvectorsBidiagU wsWU 0 wsDOrig 0 wsW 0 nn_ + + -- V-GEMM: wsVres = Vsub * W_V + do baVs <- unsafeFreezeByteArray wsVsub + baWv <- unsafeFreezeByteArray wsW + rawZeroDoubles wsVres 0 (maxN * nn_) + rawGemmKernel baVs 0 baWv 0 wsVres 0 maxN nn_ nn_ + forM_ [0..nn_-1] $ \i -> + rawCopyColumn wsVres 0 nn_ i wsVlocal 0 maxN (loL + i) maxN + + -- U-GEMM: wsUres = Usub * W_U + do baUs <- unsafeFreezeByteArray wsUsub + baWu <- unsafeFreezeByteArray wsWU + rawZeroDoubles wsUres 0 (maxN * nn_) + rawGemmKernel baUs 0 baWu 0 wsUres 0 maxN nn_ nn_ + forM_ [0..nn_-1] $ \i -> + rawCopyColumn wsUres 0 nn_ i wsUlocal 0 maxN (loL + i) maxN + + -- Write singular values = sqrt(|mu|) + forM_ [0..nn_-1] $ \i -> do + mu <- readRawD wsLam 0 i + writeRawD mbaD offD (lo + i) (sqrt (max 0 (abs mu))) + + else do + -- Partial deflation: reduced secular solve + reduced GEMM + -- Build compressed d²_nd and z_nd in wsQtemp + forM_ [0..kND-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + dpi <- readRawD wsDSort 0 pi_ + zpi <- readRawD wsZSort 0 pi_ + writeRawD wsQtemp 0 j dpi + writeRawD wsQtemp 0 (kND + j) zpi + + -- Also build compressed dOrig_nd for U-eigenvectors + forM_ [0..kND-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + dO <- readRawD wsDOrig 0 pi_ + writeRawD wsQtemp 0 (2 * kND + j) dO + + secularSolve wsLam 0 wsQtemp 0 wsQtemp kND rho kND deflTol + dcEigenvectors wsW 0 wsQtemp 0 wsQtemp kND wsLam 0 rho kND + dcEigenvectorsBidiagU wsWU 0 wsQtemp (2 * kND) wsW 0 kND + + -- Copy deflated columns to local accumulators + forM_ [kND..nn_-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + rawCopyColumn wsVsub 0 nn_ pi_ wsVlocal 0 maxN (loL + j) maxN + rawCopyColumn wsUsub 0 nn_ pi_ wsUlocal 0 maxN (loL + j) maxN + + -- Extract V_nd (maxN × kND) from non-deflated columns + forM_ [0..kND-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + rawCopyColumn wsVsub 0 nn_ pi_ wsVres 0 kND j maxN + + -- V-GEMM: wsVsub(maxN×kND) = V_nd(maxN×kND) * W_V(kND×kND) + do baVnd <- unsafeFreezeByteArray wsVres + baWv <- unsafeFreezeByteArray wsW + rawZeroDoubles wsVsub 0 (maxN * kND) + rawGemmKernel baVnd 0 baWv 0 wsVsub 0 maxN kND kND + forM_ [0..kND-1] $ \j -> + rawCopyColumn wsVsub 0 kND j wsVlocal 0 maxN (loL + j) maxN + + -- Extract U_nd (maxN × kND) from non-deflated columns + forM_ [0..kND-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + rawCopyColumn wsUsub 0 nn_ pi_ wsUres 0 kND j maxN + + -- U-GEMM: wsUsub(maxN×kND) = U_nd(maxN×kND) * W_U(kND×kND) + do baUnd <- unsafeFreezeByteArray wsUres + baWu <- unsafeFreezeByteArray wsWU + rawZeroDoubles wsUsub 0 (maxN * kND) + rawGemmKernel baUnd 0 baWu 0 wsUsub 0 maxN kND kND + forM_ [0..kND-1] $ \j -> + rawCopyColumn wsUsub 0 kND j wsUlocal 0 maxN (loL + j) maxN + + -- Write eigenvalues: non-deflated from secular, deflated from sorted d² + forM_ [0..kND-1] $ \i -> do + mu <- readRawD wsLam 0 i + writeRawD mbaD offD (lo + i) (sqrt (max 0 (abs mu))) + forM_ [kND..nn_-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + dsq <- readRawD wsDSort 0 pi_ + writeRawD mbaD offD (lo + j) (sqrt (max 0 dsq)) + + -- Run the D&C recursion + dcGo 0 (nn - 1) + + -- Final step: apply local accumulators to global U and V via GEMM + -- V[:, 0..nn-1] = V[:, 0..nn-1] * Vlocal + forM_ [0..nn-1] $ \j -> + rawCopyColumn mbaV offV nn0 j wsVsub 0 nn j nn0 + do baVs <- unsafeFreezeByteArray wsVsub + baVl <- unsafeFreezeByteArray wsVlocal + rawZeroDoubles wsVres 0 (nn0 * nn) + rawGemmKernel baVs 0 baVl 0 wsVres 0 nn0 nn nn + forM_ [0..nn-1] $ \j -> + rawCopyColumn wsVres 0 nn j mbaV offV nn0 j nn0 + + -- U[:, 0..nn-1] = U[:, 0..nn-1] * Ulocal + forM_ [0..nn-1] $ \j -> + rawCopyColumn mbaU offU mm j wsUsub 0 nn j mm + do baUs <- unsafeFreezeByteArray wsUsub + baUl <- unsafeFreezeByteArray wsUlocal + rawZeroDoubles wsUres 0 (mm * nn) + rawGemmKernel baUs 0 baUl 0 wsUres 0 mm nn nn + forM_ [0..nn-1] $ \j -> + rawCopyColumn wsUres 0 nn j mbaU offU mm j mm + +-- | Compute U-eigenvectors from V-eigenvectors and original (unsquared) d-values. +-- W_U[j,i] = dOrig[j] * W_V[j,i], then normalize each column. +dcEigenvectorsBidiagU :: MutableByteArray s -> Int -- W_U output (nn × nn) + -> MutableByteArray s -> Int -- dOrig (unsquared singular values) + -> MutableByteArray s -> Int -- W_V (V-eigenvectors, already computed) + -> Int -- nn + -> ST s () +dcEigenvectorsBidiagU mbaWU offWU mbaDOrig offDO mbaWV offWV nn = do + forM_ [0..nn-1] $ \i -> do + -- W_U[:,i] = diag(dOrig) * W_V[:,i], then normalize + norm2 <- goCol i 0 0 + let !invNorm = if norm2 > 0 then 1 / sqrt norm2 else 1 + forM_ [0..nn-1] $ \j -> do + wuji <- readRawD mbaWU offWU (j * nn + i) + writeRawD mbaWU offWU (j * nn + i) (wuji * invNorm) + where + goCol !i !j !acc + | j >= nn = pure acc + | otherwise = do + dj <- readRawD mbaDOrig offDO j + wvji <- readRawD mbaWV offWV (j * nn + i) + let !wuji = dj * wvji + writeRawD mbaWU offWU (j * nn + i) wuji + goCol i (j + 1) (acc + wuji * wuji) +{-# NOINLINE dcBidiagSVD #-} diff --git a/src/Numeric/LinearAlgebra/Massiv/Eigen/Schur.hs b/src/Numeric/LinearAlgebra/Massiv/Eigen/Schur.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Eigen/Schur.hs @@ -0,0 +1,259 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Eigen.Schur +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Real Schur decomposition via the practical QR algorithm, following +-- Golub & Van Loan, /Matrix Computations/, 4th edition (GVL4), Section 7.5, +-- pp. 393--417. +-- +-- __Theorem 7.5.1 (Real Schur Decomposition, p. 393):__ For every +-- \(A \in \mathbb{R}^{n \times n}\) there exists an orthogonal matrix \(Q\) +-- such that +-- +-- \[ +-- A = Q \, T \, Q^T +-- \] +-- +-- where \(T\) is upper /quasi/-triangular: its diagonal consists of \(1 +-- \times 1\) blocks (real eigenvalues) and \(2 \times 2\) blocks whose +-- eigenvalues are complex conjugate pairs \(\alpha \pm \beta i\). +-- +-- __Algorithm:__ The implementation follows GVL4 Algorithm 7.5.1 (Practical +-- QR Algorithm, p. 395): +-- +-- 1. Reduce \(A\) to upper Hessenberg form \(H\) via +-- "Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg". +-- 2. Apply implicit single-shift QR iterations with Givens rotations on +-- \(H\), using the /Wilkinson shift/ (eigenvalue of the trailing \(2 +-- \times 2\) block closest to \(h_{nn}\), p. 397) to accelerate +-- convergence. +-- 3. Deflate converged eigenvalues from the bottom of the active +-- Hessenberg window. +-- +-- The Wilkinson shift ensures global convergence; in practice, most +-- eigenvalues converge in only one or two iterations (GVL4, p. 397). +module Numeric.LinearAlgebra.Massiv.Eigen.Schur + ( -- * Schur decomposition (Algorithm 7.5.1) + schur + -- * Eigenvalues from Schur form + , eigenvalues + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg (hessenberg) +import Numeric.LinearAlgebra.Massiv.Orthogonal.Givens (givensRotation) + +-- | Real Schur decomposition (GVL4 Theorem 7.5.1, p. 393; Algorithm 7.5.1, +-- p. 395). +-- +-- Computes orthogonal \(Q\) and upper quasi-triangular \(T\) satisfying +-- +-- \[ +-- A = Q \, T \, Q^T +-- \] +-- +-- The matrix \(T\) has the same eigenvalues as \(A\). Its diagonal blocks +-- are either: +-- +-- * \(1 \times 1\) — corresponding to a real eigenvalue, or +-- * \(2 \times 2\) — corresponding to a pair of complex conjugate +-- eigenvalues \(\alpha \pm \beta i\). +-- +-- Internally the algorithm first reduces \(A\) to upper Hessenberg form via +-- 'Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg.hessenberg', then applies +-- implicit single-shift QR iterations using the /Wilkinson shift/ (GVL4, +-- p. 397) and Givens rotations. +-- +-- Returns @(Q, T)@. +schur :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e + -> Int -- ^ Maximum iterations + -> e -- ^ Convergence tolerance + -> (Matrix n n r e, Matrix n n r e) +schur a maxIter tol = + let nn = dimVal @n + -- Step 1: Reduce to Hessenberg form + (q0, h0) = hessenberg a + -- Step 2: QR iteration on Hessenberg matrix + (qFinal, tFinal) = qrIteration nn q0 h0 maxIter tol + in (qFinal, tFinal) + +-- | Implicit QR iteration on an upper Hessenberg matrix. +qrIteration :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Int -> Matrix n n r e -> Matrix n n r e -> Int -> e + -> (Matrix n n r e, Matrix n n r e) +qrIteration nn q h maxIter tol = go 0 q h (nn - 1) + where + go :: Int -> Matrix n n r e -> Matrix n n r e -> Int -> (Matrix n n r e, Matrix n n r e) + go iter q_ h_ p + | iter >= maxIter = (q_, h_) + | p <= 0 = (q_, h_) + | otherwise = + -- Check for convergence of h(p, p-1) + let subdiag = abs (h_ ! (p, p - 1)) + diagSum = abs (h_ ! (p - 1, p - 1)) + abs (h_ ! (p, p)) + in if subdiag <= tol * diagSum + then + -- Deflate: set subdiagonal to zero, reduce problem size + let h_new = makeMatrix @n @n @r $ \i j -> + if i == p && j == p - 1 then 0 else h_ ! (i, j) + in go iter q_ h_new (p - 1) + else + -- Apply one QR step with Wilkinson shift + let shift = wilkinsonShift (h_ ! (p-1, p-1)) (h_ ! (p-1, p)) + (h_ ! (p, p-1)) (h_ ! (p, p)) + -- Shifted QR step: H - σI = QR, H_new = RQ + σI + -- Implemented via Givens rotations on Hessenberg matrix + (q_new, h_new) = qrStepGivens q_ h_ shift p + in go (iter + 1) q_new h_new p + +-- | Wilkinson shift (GVL4, p. 397). +-- +-- Given the trailing \(2 \times 2\) block +-- +-- \[ +-- \begin{bmatrix} a & b \\ c & d \end{bmatrix} +-- \] +-- +-- the Wilkinson shift is the eigenvalue of this block that is closest to +-- \(d\) (the bottom-right entry). When the eigenvalues of the block are +-- complex the shift defaults to \(d\). +wilkinsonShift :: (Floating e, Ord e) => e -> e -> e -> e -> e +wilkinsonShift a b c d = + let trace_ = a + d + det_ = a * d - b * c + disc = trace_ * trace_ / 4 - det_ + in if disc < 0 + then d -- Complex eigenvalues; use d as shift + else + let sqrtDisc = sqrt disc + mu1 = trace_ / 2 + sqrtDisc + mu2 = trace_ / 2 - sqrtDisc + in if abs (mu1 - d) < abs (mu2 - d) then mu1 else mu2 + +-- | One QR step on Hessenberg matrix using Givens rotations. +-- H ← shift, QR factorize, then H = RQ + shift. +qrStepGivens :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Matrix n n r e -> e -> Int + -> (Matrix n n r e, Matrix n n r e) +qrStepGivens q h shift p = + let nn = dimVal @n + -- Apply shift: H ← H - σI + h_shifted = makeMatrix @n @n @r $ \i j -> + if i == j then (h ! (i, j)) - shift else h ! (i, j) + -- QR factorization via Givens rotations (only on the active part) + (rotations, r) = applyGivensQR h_shifted p + -- Form RQ + σI + h_new = formRQ r rotations shift p + -- Update Q + q_new = updateQ q rotations p + in (q_new, h_new) + +-- | Apply Givens rotations to zero out subdiagonal of Hessenberg matrix. +-- Returns list of (c, s, row_index) and the resulting R. +applyGivensQR :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Int -> ([(e, e, Int)], Matrix n n r e) +applyGivensQR h p = foldl step ([], h) [0..p-1] + where + nn = dimVal @n + step (rots, hh) k = + let (c, s) = givensRotation (hh ! (k, k)) (hh ! (k+1, k)) + hh' = applyGivensLeftSq c s k (k+1) hh + in (rots ++ [(c, s, k)], hh') + +-- | Apply Givens from left to a square matrix. +applyGivensLeftSq :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => e -> e -> Int -> Int -> Matrix n n r e -> Matrix n n r e +applyGivensLeftSq c s ri rk h = + makeMatrix @n @n @r $ \i j -> + if i == ri then + c * (h ! (ri, j)) - s * (h ! (rk, j)) + else if i == rk then + s * (h ! (ri, j)) + c * (h ! (rk, j)) + else + h ! (i, j) + +-- | Form RQ + σI from R and the Givens rotations. +formRQ :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Matrix n n r e -> [(e, e, Int)] -> e -> Int -> Matrix n n r e +formRQ r rots shift _ = + let -- Apply rotations from the right: R·G₁ᵀ·G₂ᵀ·... + rq = foldl (\mat (c, s, k) -> + applyGivensRightSq c s k (k+1) mat + ) r rots + in -- Add back shift + makeMatrix @(MatDim n) @(MatDim n) $ \i j -> + if i == j then (rq ! (i, j)) + shift else rq ! (i, j) + +type MatDim n = n -- type alias to avoid ambiguity + +-- | Apply Givens from right to a square matrix. +applyGivensRightSq :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => e -> e -> Int -> Int -> Matrix n n r e -> Matrix n n r e +applyGivensRightSq c s ci ck h = + makeMatrix @n @n @r $ \i j -> + if j == ci then + c * (h ! (i, ci)) - s * (h ! (i, ck)) + else if j == ck then + s * (h ! (i, ci)) + c * (h ! (i, ck)) + else + h ! (i, j) + +-- | Update Q by applying Givens rotations from the right. +updateQ :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Matrix n n r e -> [(e, e, Int)] -> Int -> Matrix n n r e +updateQ q rots _ = foldl (\qq (c, s, k) -> + applyGivensRightSq c s k (k+1) qq + ) q rots + +-- | Extract eigenvalues from a (quasi-)upper triangular Schur form \(T\). +-- +-- The Schur matrix \(T\) produced by 'schur' has \(1 \times 1\) and +-- \(2 \times 2\) diagonal blocks. This function walks the diagonal and +-- extracts eigenvalues: +-- +-- * A \(1 \times 1\) block \([t_{ii}]\) yields the real eigenvalue +-- \(\lambda = t_{ii}\). +-- * A \(2 \times 2\) block +-- \(\bigl[\begin{smallmatrix} a & b \\ c & d \end{smallmatrix}\bigr]\) +-- yields eigenvalues \(\tfrac{a + d}{2} \pm \sqrt{\tfrac{(a+d)^2}{4} - +-- (ad - bc)}\). When the discriminant is negative (complex conjugate +-- pair) only the real part \(\tfrac{a + d}{2}\) is returned for each +-- eigenvalue, since this module operates over real scalars. +-- +-- See GVL4 Section 7.5 for the definition of the real Schur form. +eigenvalues :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> [e] +eigenvalues t = + let nn = dimVal @n + in go 0 + where + nn = dimVal @n + go i + | i >= nn = [] + | i == nn - 1 = [t ! (i, i)] -- Last 1×1 block + | abs (t ! (i+1, i)) < 1e-12 * (abs (t ! (i, i)) + abs (t ! (i+1, i+1))) = + -- 1×1 block + t ! (i, i) : go (i + 1) + | otherwise = + -- 2×2 block: eigenvalues of [[a,b],[c,d]] + let a = t ! (i, i) + b = t ! (i, i+1) + c = t ! (i+1, i) + d = t ! (i+1, i+1) + tr = a + d + det_ = a * d - b * c + disc = tr * tr / 4 - det_ + in if disc >= 0 + then (tr / 2 + sqrt disc) : (tr / 2 - sqrt disc) : go (i + 2) + else tr / 2 : tr / 2 : go (i + 2) -- Complex pair, return real parts diff --git a/src/Numeric/LinearAlgebra/Massiv/Eigen/Symmetric.hs b/src/Numeric/LinearAlgebra/Massiv/Eigen/Symmetric.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Eigen/Symmetric.hs @@ -0,0 +1,1899 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Eigen.Symmetric +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Eigenvalue algorithms specialised to real symmetric matrices, following +-- Golub & Van Loan, /Matrix Computations/, 4th edition (GVL4), Chapter 8, +-- pp. 449--512. +module Numeric.LinearAlgebra.Massiv.Eigen.Symmetric + ( tridiagonalize + , tridiagonalizeP + , symmetricEigen + , symmetricEigenP + , symmetricEigenPPar + , symmetricEigenPDC + , jacobiEigen + -- * D&C secular equation infrastructure (for bidiagonal SVD reuse) + , secularSolve + , secularSolveOne + , deflatePartition + , dcEigenvectors + , secularFuncSplit + , secularFuncAndDeriv + , sumZSq + , farPoleSum + , farPoleSumSkip + , readRawD + , writeRawD + , readRawI + , writeRawI + , indexRawD + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Ix1, unwrapByteArray, unwrapByteArrayOffset, unwrapMutableByteArray, unwrapMutableByteArrayOffset) +import GHC.TypeNats (KnownNat) +import Control.Monad (when, forM_) +import Control.Monad.ST (ST, stToIO) +import Control.Concurrent (forkIO, newEmptyMVar, putMVar, takeMVar) +import System.IO.Unsafe (unsafePerformIO) + +import GHC.Exts +import GHC.ST (ST(..)) +import Data.Primitive.ByteArray (ByteArray(..), MutableByteArray(..), newByteArray, unsafeFreezeByteArray) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Orthogonal.Givens (givensRotation) +import Numeric.LinearAlgebra.Massiv.Internal.Kernel + ( rawMutApplyGivensColumns + , rawMutApplyGivensColumnsCM + , rawMutSumSqColumn + , rawMutSymMatvecSub + , rawMutSymRank2Update + , rawMutTridiagQAccum + , rawGemmKernel + , rawTransposeToColMajor + , rawTransposeFromColMajor + , rawZeroDoubles + , rawCopyDoubles + , rawNegateDoubles + , rawCopyColumn + ) + +-- | Reduce a symmetric matrix to tridiagonal form (GVL4 Algorithm 8.3.1). +tridiagonalize :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e + -> (Matrix n n r e, Vector n r e, Vector n r e) +tridiagonalize a = + let nn = dimVal @n + + -- Phase 1: In-place tridiagonalisation via symmetric rank-2 updates. + (betaList, tArr) = M.withMArrayST (unMatrix a) $ \mt -> do + betas <- mapM (tridiagStep mt nn) [0..nn-3] + pure betas + + -- Phase 2: Accumulate Q from stored Householder vectors. + qMat = createMatrix @n @n @r $ \mq -> do + forM_ [0..nn-1] $ \i -> forM_ [0..nn-1] $ \j -> + M.write_ mq (i :. j) (if i == j then 1 else 0) + -- Forward accumulation: Q <- Q · H_k for k = 0..n-3 + forM_ (zip [0..] betaList) $ \(k, beta_k) -> + when (beta_k /= 0) $ + forM_ [0..nn-1] $ \i -> do + qik1 <- M.readM mq (i :. (k+1)) + rest <- sumQV mq tArr i (k+1) nn k + let wi = beta_k * (qik1 + rest) + M.write_ mq (i :. (k+1)) (qik1 - wi) + forM_ [k+2..nn-1] $ \l -> do + let vl = M.index' tArr (l :. k) + qil <- M.readM mq (i :. l) + M.write_ mq (i :. l) (qil - wi * vl) + + diag_ = makeVector @n @r $ \i -> M.index' tArr (i :. i) + subdiag = makeVector @n @r $ \i -> + if i < nn - 1 then M.index' tArr ((i+1) :. i) else 0 + + in (qMat, diag_, subdiag) + +-- | One step of Householder tridiagonalisation. +tridiagStep :: (M.Manifest r e, Floating e, Ord e) + => M.MArray s r Ix2 e -> Int -> Int -> ST s e +tridiagStep mt nn k = do + x0 <- M.readM mt ((k+1) :. k) + sigma <- sumSqBelow mt (k+1) nn k + if sigma == 0 && x0 >= 0 + then pure 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + -- Build v as a list: v(k+1)=1, v(i)=T(i,k)/v0 for i>k+1 + vList <- mapM (\i -> do + tik <- M.readM mt (i :. k) + pure (tik / v0) + ) [k+2..nn-1] + let fullV = 1 : vList -- indices k+1, k+2, ..., n-1 + -- p = beta * T * v (rows k+1..n-1) + pList <- mapM (\i -> do + s <- dotTV mt i fullV (k+1) nn + pure (beta * s) + ) [k+1..nn-1] + let ptv = sum $ zipWith (*) pList fullV + alpha_ = beta * ptv / 2 + wList = zipWith (\pi_ vi -> pi_ - alpha_ * vi) pList fullV + -- Symmetric rank-2 update: T(i,j) -= v(i)*w(j) + w(i)*v(j) + forM_ (zip3 [k+1..nn-1] fullV wList) $ \(i, vi, wi) -> + forM_ (zip3 [k+1..nn-1] fullV wList) $ \(j, vj, wj) -> do + tij <- M.readM mt (i :. j) + M.write_ mt (i :. j) (tij - vi * wj - wi * vj) + -- Store Householder vector in below-subdiagonal of column k + forM_ (zip [k+2..nn-1] vList) $ \(i, vi) -> + M.write_ mt (i :. k) vi + -- Set subdiagonal + M.write_ mt ((k+1) :. k) mu + M.write_ mt (k :. (k+1)) mu + pure beta + +-- Helpers for tridiagonalize +sumSqBelow :: (M.Manifest r e, Num e) => M.MArray s r Ix2 e -> Int -> Int -> Int -> ST s e +sumSqBelow mt start end col = go (start + 1) 0 + where go i !acc | i >= end = pure acc + | otherwise = do v <- M.readM mt (i :. col); go (i+1) (acc + v*v) + +dotTV :: (M.Manifest r e, Num e) => M.MArray s r Ix2 e -> Int -> [e] -> Int -> Int -> ST s e +dotTV mt i vList start end = go start vList 0 + where go _ [] !acc = pure acc + go j (v:vs) !acc | j >= end = pure acc + | otherwise = do t <- M.readM mt (i :. j); go (j+1) vs (acc + t*v) + +sumQV :: (M.Manifest r1 e, M.Manifest r2 e, Num e) + => M.MArray s r1 Ix2 e -> M.Array r2 Ix2 e -> Int -> Int -> Int -> Int -> ST s e +sumQV mq tArr row start end col = go (start + 1) 0 + where go l !acc | l >= end = pure acc + | otherwise = do + q <- M.readM mq (row :. l) + let v = M.index' tArr (l :. col) + go (l+1) (acc + q*v) + +-- | Raw-primop tridiagonalisation specialised for @P Double@. +-- Two-phase: (1) in-place Householder via raw ByteArray# kernels, +-- (2) Q accumulation using rawMutTridiagQAccum. +tridiagonalizeP :: forall n. KnownNat n + => Matrix n n M.P Double + -> (Matrix n n M.P Double, Vector n M.P Double, Vector n M.P Double) +tridiagonalizeP a = + let nn = dimVal @n + + -- Phase 1: In-place Householder tridiagonalisation + -- For n < panelCrossover: per-column Level-2 (rank-2 update per step) + -- For n >= panelCrossover: DLATRD-style panel factorisation (Level-3 SYR2K) + panelCrossover = 64 + (betaList, tArr) = M.withMArrayST (unMatrix a) $ \mt -> do + let !mbaT = unwrapMutableByteArray mt + !offT = unwrapMutableByteArrayOffset mt + mbaV <- newByteArray (nn * 8) + mbaP <- newByteArray (nn * 8) + mbaW <- newByteArray (nn * 8) + if nn < panelCrossover + then do + betas <- mapM (\k -> tridiagStepP mbaT offT nn mbaV mbaP mbaW k) [0..nn-3] + pure betas + else do + let !nb = min 64 (max 16 (nn `div` 3)) + !numRef = nn - 2 -- number of Householder reflectors + -- V_panel (nn × nb) and W_panel (nn × nb) for deferred rank-2 updates + mbaVp <- newByteArray (nn * nb * 8) + mbaWp <- newByteArray (nn * nb * 8) + -- Temporary for GEMM-based trailing update + mbaTemp <- newByteArray (nn * nb * 8) + -- Pre-allocate workspace for panelTridiagP (avoids per-panel allocation) + wsHvSave <- newByteArray (nb * nn * 8) + wsVr <- newByteArray (nn * nb * 8) + wsWr <- newByteArray (nn * nb * 8) + wsNWrT <- newByteArray (nb * nn * 8) + wsNVrT <- newByteArray (nb * nn * 8) + wsRem <- newByteArray (nn * nn * 8) + let go !k0 !accBetas + | k0 > numRef - 1 = pure (reverse accBetas) + | otherwise = do + let !bs = min nb (numRef - k0) + panelBetas <- panelTridiagP mbaT offT nn mbaV mbaP mbaW + mbaVp mbaWp mbaTemp + wsHvSave wsVr wsWr wsNWrT wsNVrT wsRem + k0 bs + go (k0 + bs) (reverse panelBetas ++ accBetas) + go 0 [] + + -- Get underlying ByteArray from frozen T for Q accumulation + !tBA = unwrapByteArray tArr + !tOff = unwrapByteArrayOffset tArr + + -- Phase 2: Q accumulation. + -- For n < 200: per-row Householder updates (minimal work, avoids GEMM overhead). + -- For n >= 200: blocked WY with Level-3 GEMM (better cache/SIMD utilization). + qMat = createMatrix @n @n @M.P $ \mq -> do + let !mbaQ = unwrapMutableByteArray mq + !offQ = unwrapMutableByteArrayOffset mq + -- Set Q = I (SIMD zero + diagonal ones) + rawZeroDoubles mbaQ offQ (nn * nn) + forM_ [0..nn-1] $ \i -> writeRawD mbaQ offQ (i*nn+i) 1 + if nn < 128 + then + -- Per-row approach: Q <- Q · H_k for k = 0..n-3 + forM_ (zip [0..] betaList) $ \(k, beta_k) -> + when (beta_k /= 0) $ + forM_ [0..nn-1] $ \row -> + rawMutTridiagQAccum mbaQ offQ nn tBA tOff nn beta_k (k+1) k nn row + else do + -- Blocked WY approach: Q <- Q * (I - Y * T * Y^T) per block + let !numRef = nn - 2 + !nb = min 48 numRef + mbaBetas <- newByteArray (numRef * 8) + forM_ (zip [0..] betaList) $ \(i, b) -> writeRawD mbaBetas 0 i b + mbaY <- newByteArray (nn * nb * 8) + mbaTf <- newByteArray (nb * nb * 8) + mbaW1 <- newByteArray (nn * nb * 8) + mbaW2 <- newByteArray (nn * nb * 8) + mbaYT <- newByteArray (nb * nn * 8) + mbaG <- newByteArray (nb * nb * 8) -- Gram matrix Y^T Y + + forM_ [0, nb .. numRef - 1] $ \k0 -> do + let !bs = min nb (numRef - k0) + + -- Pack Y (n × bs) from stored Householder vectors + rawZeroDoubles mbaY 0 (nn * bs) + forM_ [0..bs-1] $ \j -> do + let !k = k0 + j + writeRawD mbaY 0 ((k+1) * bs + j) 1.0 + forM_ [k+2..nn-1] $ \l -> + writeRawD mbaY 0 (l * bs + j) (indexRawD tBA tOff (l * nn + k)) + + -- Transpose Y → Y^T (bs × n) early: reused for T factor and final GEMM + forM_ [0..nn-1] $ \row -> + forM_ [0..bs-1] $ \col -> + writeRawD mbaYT 0 (col * nn + row) 0 + forM_ [0..bs-1] $ \j -> do + let !k = k0 + j + writeRawD mbaYT 0 (j * nn + (k+1)) 1.0 + forM_ [k+2..nn-1] $ \l -> + writeRawD mbaYT 0 (j * nn + l) (indexRawD tBA tOff (l * nn + k)) + + -- Freeze Y and Y^T for GEMM use + baY <- unsafeFreezeByteArray mbaY + baYT <- unsafeFreezeByteArray mbaYT + + -- Compute G = Y^T × Y (bs × bs) via GEMM for T factor dot products + rawZeroDoubles mbaG 0 (bs * bs) + rawGemmKernel baYT 0 baY 0 mbaG 0 bs nn bs + + -- Build T factor (bs × bs upper-triangular) using precomputed G + rawZeroDoubles mbaTf 0 (bs * bs) + forM_ [0..bs-1] $ \j -> do + betaj <- readRawD mbaBetas 0 (k0 + j) + writeRawD mbaTf 0 (j * bs + j) betaj + when (j > 0 && betaj /= 0) $ do + -- Read G[i,j] = Y[:,i]^T Y[:,j] for all i < j + forM_ [0..j-1] $ \i -> do + g_ij <- readRawD mbaG 0 (i * bs + j) + writeRawD mbaW1 0 i g_ij + -- Triangular solve: T[i,j] = -betaj * Σ_l T[i,l] * G[l,j] + forM_ [0..j-1] $ \i -> do + let triLoop !l !acc + | l >= j = pure acc + | otherwise = do + til <- readRawD mbaTf 0 (i * bs + l) + dl <- readRawD mbaW1 0 l + triLoop (l+1) (acc + til * dl) + z <- triLoop i 0 + writeRawD mbaTf 0 (i * bs + j) (negate betaj * z) + + -- W1 = Q · Y (GEMM n×n * n×bs → n×bs) + baQ <- unsafeFreezeByteArray mbaQ + rawZeroDoubles mbaW1 0 (nn * bs) + rawGemmKernel baQ offQ baY 0 mbaW1 0 nn nn bs + + -- W2 = W1 · T (GEMM n×bs * bs×bs → n×bs) + baW1 <- unsafeFreezeByteArray mbaW1 + baTf <- unsafeFreezeByteArray mbaTf + rawZeroDoubles mbaW2 0 (nn * bs) + rawGemmKernel baW1 0 baTf 0 mbaW2 0 nn bs bs + + -- Negate W2 in-place (SIMD) + rawNegateDoubles mbaW2 0 (nn * bs) + + -- Q += (-W2) · Y^T (GEMM n×bs * bs×n → n×n) — reuses baYT + baNW2 <- unsafeFreezeByteArray mbaW2 + rawGemmKernel baNW2 0 baYT 0 mbaQ offQ nn bs nn + + -- Read diagonal and subdiagonal from frozen T + diag_ = makeVector @n @M.P $ \i -> M.index' tArr (i :. i) + subdiag = makeVector @n @M.P $ \i -> + if i < nn - 1 then M.index' tArr ((i+1) :. i) else 0 + + in (qMat, diag_, subdiag) +{-# NOINLINE tridiagonalizeP #-} + +-- | One step of raw-primop Householder tridiagonalisation. +tridiagStepP :: MutableByteArray s -> Int -> Int + -> MutableByteArray s -> MutableByteArray s -> MutableByteArray s + -> Int -> ST s Double +tridiagStepP mbaT offT nn mbaV mbaP mbaW k = do + -- 1. Read x0 = T[k+1,k] + x0 <- readRawD mbaT offT ((k+1)*nn + k) + -- 2. Compute sigma = Σ T[i,k]^2 for i=k+2..nn-1 + sigma <- rawMutSumSqColumn mbaT offT nn (k+2) nn k + if sigma == 0 && x0 >= 0 + then pure 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + subSize = nn - k - 1 + + -- 3. Build v in mbaV: v[0]=1, v[i]=T[k+1+i,k]/v0 + writeRawD mbaV 0 0 1.0 + forM_ [1..subSize-1] $ \i -> do + tik <- readRawD mbaT offT ((k+1+i)*nn + k) + writeRawD mbaV 0 i (tik / v0) + + -- 4. p = beta * T_sub * v + rawMutSymMatvecSub mbaT offT nn mbaV 0 mbaP 0 (k+1) nn + forM_ [0..subSize-1] $ \i -> do + pi_ <- readRawD mbaP 0 i + writeRawD mbaP 0 i (beta * pi_) + + -- 5. Dot product p^T v + ptv <- mutDotVec mbaP 0 mbaV 0 subSize + let alpha_ = beta * ptv / 2 + + -- 6. w = p - alpha*v + forM_ [0..subSize-1] $ \i -> do + pi_ <- readRawD mbaP 0 i + vi <- readRawD mbaV 0 i + writeRawD mbaW 0 i (pi_ - alpha_ * vi) + + -- 7. Rank-2 update: T -= vw^T + wv^T + rawMutSymRank2Update mbaT offT nn mbaV 0 mbaW 0 (k+1) nn + + -- 8. Store Householder vector in column k subdiagonal + forM_ [1..subSize-1] $ \i -> do + vi <- readRawD mbaV 0 i + writeRawD mbaT offT ((k+1+i)*nn + k) vi + + -- 9. Set subdiagonal element + writeRawD mbaT offT ((k+1)*nn + k) mu + writeRawD mbaT offT (k*nn + (k+1)) mu + + pure beta + +-- | DLATRD-style panel tridiagonalisation. +-- Processes columns k0..k0+bs-1, building V_panel and W_panel matrices +-- that represent the deferred rank-2 updates. After processing all columns +-- in the panel, applies a single Level-3 SYR2K trailing update. +-- +-- Within the panel, column k of T is corrected for deferred updates: +-- T[:,k] -= V_panel * W_panel[k,:] + W_panel * V_panel[k,:] +-- before computing the Householder reflector. +-- +-- Returns the list of beta values for the panel columns. +panelTridiagP :: MutableByteArray s -> Int -> Int -- T matrix, offset, n + -> MutableByteArray s -> MutableByteArray s -> MutableByteArray s -- v, p, w temps + -> MutableByteArray s -> MutableByteArray s -> MutableByteArray s -- Vp, Wp, temp + -> MutableByteArray s -> MutableByteArray s -> MutableByteArray s -- wsHvSave, wsVr, wsWr + -> MutableByteArray s -> MutableByteArray s -> MutableByteArray s -- wsNWrT, wsNVrT, wsRem + -> Int -> Int -- k0, bs (panel start, panel size) + -> ST s [Double] +panelTridiagP mbaT offT nn mbaV mbaP mbaW mbaVp mbaWp _mbaTemp + mbaHvSave mbaVr mbaWr mbaNWrT mbaNVrT mbaRem k0 bs = do + -- DLATRD-style: NO rank-2 updates to T within the panel. + -- All corrections computed from V_panel, W_panel. + -- After the panel, apply SYR2K to the full remaining submatrix. + betas <- go 0 [] + + -- After the panel: apply accumulated rank-2 update to the full remaining + -- submatrix T[k0+1:nn, k0+1:nn]. This includes both within-panel diagonal + -- entries and the trailing submatrix. + -- + -- We must save/restore Householder vectors in columns k0..k0+bs-1 + -- because the SYR2K will overwrite them. + let !remStart = k0 + 1 + !remSize = nn - remStart + when (remSize > 0 && bs > 0) $ do + -- Save Householder vectors from T columns k0..k0+bs-1 + -- These are T[i, k] for i > k+1, k in [k0..k0+bs-1] + -- Also save subdiagonal entries T[k+1, k] = mu + forM_ [0..bs-1] $ \l -> do + let !k = k0 + l + !startRow = k + 1 + forM_ [startRow..nn-1] $ \i -> do + val <- readRawD mbaT offT (i * nn + k) + writeRawD mbaHvSave 0 (l * nn + i) val + -- Also save T[k, k+1] (the upper subdiagonal) + when (k + 1 < nn) $ do + val <- readRawD mbaT offT (k * nn + (k + 1)) + writeRawD mbaHvSave 0 (l * nn + k) val -- reuse slot k < startRow + + -- Build contiguous V_rem (remSize × bs) and W_rem (remSize × bs) + -- V_panel and W_panel have stride bs, so V_rem is a contiguous subblock + rawCopyDoubles mbaVr 0 mbaVp (remStart * bs) (remSize * bs) + rawCopyDoubles mbaWr 0 mbaWp (remStart * bs) (remSize * bs) + + -- Build -W_rem^T and -V_rem^T (bs × remSize) via transpose + negate + forM_ [0..remSize-1] $ \i -> + forM_ [0..bs-1] $ \j -> do + readRawD mbaWr 0 (i * bs + j) >>= writeRawD mbaNWrT 0 (j * remSize + i) + readRawD mbaVr 0 (i * bs + j) >>= writeRawD mbaNVrT 0 (j * remSize + i) + rawNegateDoubles mbaNWrT 0 (bs * remSize) + rawNegateDoubles mbaNVrT 0 (bs * remSize) + + -- Copy T_rem to contiguous temp (row-by-row bulk copy) + forM_ [0..remSize-1] $ \i -> + rawCopyDoubles mbaRem (i * remSize) mbaT (offT + (remStart + i) * nn + remStart) remSize + + -- GEMM: rem += V_rem * (-W_rem^T) + W_rem * (-V_rem^T) + baVr <- unsafeFreezeByteArray mbaVr + baNWrT <- unsafeFreezeByteArray mbaNWrT + rawGemmKernel baVr 0 baNWrT 0 mbaRem 0 remSize bs remSize + baWr <- unsafeFreezeByteArray mbaWr + baNVrT <- unsafeFreezeByteArray mbaNVrT + rawGemmKernel baWr 0 baNVrT 0 mbaRem 0 remSize bs remSize + + -- Copy back to T (row-by-row bulk copy) + forM_ [0..remSize-1] $ \i -> + rawCopyDoubles mbaT (offT + (remStart + i) * nn + remStart) mbaRem (i * remSize) remSize + + -- Restore saved Householder vectors and subdiagonal entries + forM_ [0..bs-1] $ \l -> do + let !k = k0 + l + !startRow = k + 1 + forM_ [startRow..nn-1] $ \i -> do + val <- readRawD mbaHvSave 0 (l * nn + i) + writeRawD mbaT offT (i * nn + k) val + when (k + 1 < nn) $ do + val <- readRawD mbaHvSave 0 (l * nn + k) + writeRawD mbaT offT (k * nn + (k + 1)) val + + pure betas + where + go !j !acc + | j >= bs = pure (reverse acc) + | otherwise = do + let !k = k0 + j + !subSize = nn - k - 1 + + -- Step 1: Read corrected column. T is ORIGINAL (no rank-2 updates applied). + -- corrected_col[i] = T[i+k+1, k] - Σ_l (V[i+k+1,l]*W[k,l] + W[i+k+1,l]*V[k,l]) + forM_ [0..subSize-1] $ \i -> do + tik <- readRawD mbaT offT ((k+1+i)*nn + k) + if j == 0 + then writeRawD mbaP 0 i tik + else do + let corrLoop !l !accC + | l >= j = pure accC + | otherwise = do + vp_il <- readRawD mbaVp 0 ((k+1+i) * bs + l) + wp_kl <- readRawD mbaWp 0 (k * bs + l) + wp_il <- readRawD mbaWp 0 ((k+1+i) * bs + l) + vp_kl <- readRawD mbaVp 0 (k * bs + l) + corrLoop (l+1) (accC + vp_il * wp_kl + wp_il * vp_kl) + corr <- corrLoop 0 0 + writeRawD mbaP 0 i (tik - corr) + + -- Step 2: Householder from corrected column + x0 <- readRawD mbaP 0 0 + sigma <- do + let sigLoop !i !acc_ + | i >= subSize = pure acc_ + | otherwise = do + ci <- readRawD mbaP 0 i + sigLoop (i+1) (acc_ + ci * ci) + sigLoop 1 0 + if sigma == 0 && x0 >= 0 + then do + forM_ [0..nn-1] $ \i -> do + writeRawD mbaVp 0 (i * bs + j) 0 + writeRawD mbaWp 0 (i * bs + j) 0 + go (j+1) (0 : acc) + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + + writeRawD mbaV 0 0 1.0 + forM_ [1..subSize-1] $ \i -> do + ci <- readRawD mbaP 0 i + writeRawD mbaV 0 i (ci / v0) + + -- Step 3: p = beta * T * v (using ORIGINAL T, then correct via V,W) + rawMutSymMatvecSub mbaT offT nn mbaV 0 mbaP 0 (k+1) nn + forM_ [0..subSize-1] $ \i -> do + pi_ <- readRawD mbaP 0 i + writeRawD mbaP 0 i (beta * pi_) + + -- Step 4: Full V,W correction. + -- p -= beta * (V_sub*(W_sub^T*v) + W_sub*(V_sub^T*v)) + -- where V_sub = V_panel[k+1:nn-1, 0:j-1], W_sub = W_panel[k+1:nn-1, 0:j-1] + when (j > 0) $ do + forM_ [0..j-1] $ \l -> do + let dotW !idx !accW + | idx >= subSize = pure accW + | otherwise = do + wp <- readRawD mbaWp 0 ((k+1+idx) * bs + l) + vi <- readRawD mbaV 0 idx + dotW (idx+1) (accW + wp * vi) + z1 <- dotW 0 0 + + let dotV !idx !accV + | idx >= subSize = pure accV + | otherwise = do + vp <- readRawD mbaVp 0 ((k+1+idx) * bs + l) + vi <- readRawD mbaV 0 idx + dotV (idx+1) (accV + vp * vi) + z2 <- dotV 0 0 + + forM_ [0..subSize-1] $ \i -> do + vp_il <- readRawD mbaVp 0 ((k+1+i) * bs + l) + wp_il <- readRawD mbaWp 0 ((k+1+i) * bs + l) + pi_ <- readRawD mbaP 0 i + writeRawD mbaP 0 i (pi_ - beta * (vp_il * z1 + wp_il * z2)) + + -- Step 5: w = p - alpha*v + ptv <- mutDotVec mbaP 0 mbaV 0 subSize + let alpha_ = beta * ptv / 2 + forM_ [0..subSize-1] $ \i -> do + pi_ <- readRawD mbaP 0 i + vi <- readRawD mbaV 0 i + writeRawD mbaW 0 i (pi_ - alpha_ * vi) + + -- Step 6: Store v,w in panels + forM_ [0..k] $ \i -> do + writeRawD mbaVp 0 (i * bs + j) 0 + writeRawD mbaWp 0 (i * bs + j) 0 + forM_ [0..subSize-1] $ \i -> do + readRawD mbaV 0 i >>= writeRawD mbaVp 0 ((k+1+i) * bs + j) + readRawD mbaW 0 i >>= writeRawD mbaWp 0 ((k+1+i) * bs + j) + + -- Step 7: Store Householder vector in T (for Q accumulation) + forM_ [1..subSize-1] $ \i -> + readRawD mbaV 0 i >>= writeRawD mbaT offT ((k+1+i)*nn + k) + + -- Step 8: Set subdiagonal + writeRawD mbaT offT ((k+1)*nn + k) mu + writeRawD mbaT offT (k*nn + (k+1)) mu + + go (j+1) (beta : acc) + +-- | Dot product of two mutable vectors. +mutDotVec :: MutableByteArray s -> Int -> MutableByteArray s -> Int -> Int -> ST s Double +mutDotVec mbaA offA mbaB offB n = go 0 0 + where + go !i !acc + | i >= n = pure acc + | otherwise = do + ai <- readRawD mbaA offA i + bi <- readRawD mbaB offB i + go (i+1) (acc + ai * bi) + +-- | Symmetric eigenvalue decomposition (GVL4 Algorithm 8.3.3). +symmetricEigen :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Int -> e -> (Vector n r e, Matrix n n r e) +symmetricEigen a maxIter tol = + let nn = dimVal @n + (q0, diag_, subdiag) = tridiagonalize a + (dArr, qArr) = M.withMArrayST (unMatrix q0) $ \mq -> do + md <- M.thawS (unVector diag_) + msd <- M.thawS (unVector subdiag) + tridiagQRLoop md msd mq nn maxIter tol + dFrozen <- M.freezeS md + pure (MkVector dFrozen) + in (dArr, MkMatrix qArr) + +-- | In-place QR iteration on tridiagonal (d, sd) with mutable Q. +-- Uses both top and bottom deflation to shrink the active range [lo..hi], +-- effectively achieving divide-and-conquer behaviour. +tridiagQRLoop :: (M.Manifest r e, Floating e, Ord e) + => M.MArray s r Ix1 e -> M.MArray s r Ix1 e -> M.MArray s r Ix2 e + -> Int -> Int -> e -> ST s () +tridiagQRLoop md msd mq nn maxIter tol = go 0 0 (nn - 1) + where + go !iter !lo !hi + | iter >= maxIter = pure () + | lo >= hi = pure () + | otherwise = do + -- Bottom deflation + sdhi <- M.readM msd (hi - 1) + dhi1 <- M.readM md (hi - 1) + dhi <- M.readM md hi + if abs sdhi <= tol * (abs dhi1 + abs dhi) + then do + M.write_ msd (hi - 1) 0 + go iter lo (hi - 1) + else do + -- Top deflation + sdlo <- M.readM msd lo + dlo <- M.readM md lo + dlo1 <- M.readM md (lo + 1) + if abs sdlo <= tol * (abs dlo + abs dlo1) + then do + M.write_ msd lo 0 + go iter (lo + 1) hi + else do + -- Interior deflation: find split point + split <- findSplit md msd lo hi tol + case split of + Just q -> do + -- Split into two subproblems [lo..q] and [q+1..hi] + M.write_ msd q 0 + go iter lo q + go iter (q + 1) hi + Nothing -> do + -- No split found: apply QR step on [lo..hi] + let sp1 = sdhi + delta = (dhi1 - dhi) / 2 + sgn = if delta >= 0 then 1 else -1 + shift = dhi - sp1*sp1 / (delta + sgn * sqrt (delta*delta + sp1*sp1)) + implicitQRStepInPlace md msd mq nn shift lo hi + go (iter + 1) lo hi + +-- | Find an interior split point where the subdiagonal is negligible. +findSplit :: (M.Manifest r e, Floating e, Ord e) + => M.MArray s r Ix1 e -> M.MArray s r Ix1 e -> Int -> Int -> e -> ST s (Maybe Int) +findSplit md msd lo hi tol = scan (lo + 1) + where + scan q + | q >= hi - 1 = pure Nothing + | otherwise = do + sdq <- M.readM msd q + dq <- M.readM md q + dq1 <- M.readM md (q + 1) + if abs sdq <= tol * (abs dq + abs dq1) + then pure (Just q) + else scan (q + 1) + +-- | One implicit symmetric QR step via bulge-chasing Givens rotations. +-- Operates on the active sub-range [lo..hi] of the tridiagonal. +implicitQRStepInPlace :: (M.Manifest r e, Floating e, Ord e) + => M.MArray s r Ix1 e -> M.MArray s r Ix1 e -> M.MArray s r Ix2 e + -> Int -> e -> Int -> Int -> ST s () +implicitQRStepInPlace md msd mq nn shift lo hi = do + dlo <- M.readM md lo + sdlo <- M.readM msd lo + chase lo (dlo - shift) sdlo + where + chase k x z = do + let (c, s) = givensRotation x z + when (k > lo) $ + M.write_ msd (k-1) (c * x - s * z) + dk <- M.readM md k + ek <- M.readM msd k + dk1 <- M.readM md (k+1) + M.write_ md k (c*c*dk - 2*c*s*ek + s*s*dk1) + M.write_ md (k+1) (s*s*dk + 2*c*s*ek + c*c*dk1) + M.write_ msd k (c*s*(dk - dk1) + (c*c - s*s)*ek) + applyGivensRightQ mq c s k (k+1) nn + if k + 1 < hi + then do + ek1 <- M.readM msd (k+1) + let z' = -s * ek1 + M.write_ msd (k+1) (c * ek1) + ek_new <- M.readM msd k + chase (k+1) ek_new z' + else pure () + +-- | Apply Givens rotation from the right to Q: Q <- Q · G(ci, ck) +-- For P Double, uses raw ByteArray# primops; generic fallback otherwise. +applyGivensRightQ :: (M.Manifest r e, Num e) + => M.MArray s r Ix2 e -> e -> e -> Int -> Int -> Int -> ST s () +applyGivensRightQ mq c s ci ck nn = + forM_ [0..nn-1] $ \row -> do + qrc <- M.readM mq (row :. ci) + qrk <- M.readM mq (row :. ck) + M.write_ mq (row :. ci) (c * qrc - s * qrk) + M.write_ mq (row :. ck) (s * qrc + c * qrk) + +-- | Specialised symmetric eigenvalue decomposition for @P Double@. +-- Uses raw ByteArray# primops for the entire QR iteration, including +-- diagonal/subdiagonal reads and writes plus Givens rotation on Q. +symmetricEigenP :: forall n. KnownNat n + => Matrix n n M.P Double -> Int -> Double -> (Vector n M.P Double, Matrix n n M.P Double) +symmetricEigenP a maxIter tol + | dimVal @n >= dcCrossover = symmetricEigenPDC a tol + | otherwise = + let nn = dimVal @n + (q0, diag_, subdiag) = tridiagonalizeP a + (dArr, qArr) = M.withMArrayST (unMatrix q0) $ \mq -> do + md <- M.thawS (unVector diag_) + msd <- M.thawS (unVector subdiag) + let !mbaD = unwrapMutableByteArray md + !offD = unwrapMutableByteArrayOffset md + !mbaSD = unwrapMutableByteArray msd + !offSD = unwrapMutableByteArrayOffset msd + !mbaQ = unwrapMutableByteArray mq + !offQ = unwrapMutableByteArrayOffset mq + -- For small n: use row-major QR loop (avoids two O(n^2) transposes) + -- For large n: column-major Q for SIMD Givens rotations + if nn < 100 + then rawTridiagQRLoop mbaD offD mbaSD offSD mbaQ offQ nn maxIter tol + else do + mbaQcm <- newByteArray (nn * nn * 8) + rawTransposeToColMajor mbaQ offQ mbaQcm 0 nn + rawTridiagQRLoopCM mbaD offD mbaSD offSD mbaQcm 0 nn maxIter tol + rawTransposeFromColMajor mbaQcm 0 mbaQ offQ nn + dFrozen <- M.freezeS md + pure (MkVector dFrozen) + in (dArr, MkMatrix qArr) + where dcCrossover = 1000 -- D&C merge GEMM overhead doesn't amortise below ~1000 +{-# NOINLINE symmetricEigenP #-} + +-- | Parallel specialised symmetric eigenvalue decomposition for @P Double@. +-- Uses raw-primop tridiagonalisation and forks independent sub-problems +-- when the QR loop finds a split point. +symmetricEigenPPar :: forall n. KnownNat n + => Matrix n n M.P Double -> Int -> Double + -> (Vector n M.P Double, Matrix n n M.P Double) +symmetricEigenPPar a maxIter tol = unsafePerformIO $ do + let nn = dimVal @n + (q0, diag_, subdiag) = tridiagonalizeP a + -- Thaw into IO (s = RealWorld) for parallel QR iteration + mq <- M.thawS (unMatrix q0) + md <- M.thawS (unVector diag_) + msd <- M.thawS (unVector subdiag) + let !mbaD = unwrapMutableByteArray md + !offD = unwrapMutableByteArrayOffset md + !mbaSD = unwrapMutableByteArray msd + !offSD = unwrapMutableByteArrayOffset msd + !mbaQ = unwrapMutableByteArray mq + !offQ = unwrapMutableByteArrayOffset mq + -- Transpose Q to column-major for SIMD Givens, run parallel QR, transpose back + mbaQcm <- stToIO $ newByteArray (nn * nn * 8) + stToIO $ rawTransposeToColMajor mbaQ offQ mbaQcm 0 nn + rawTridiagQRLoopParCM mbaD offD mbaSD offSD mbaQcm 0 nn maxIter tol + stToIO $ rawTransposeFromColMajor mbaQcm 0 mbaQ offQ nn + dFrozen <- M.freezeS md + qFrozen <- M.freezeS mq + pure (MkVector dFrozen, MkMatrix qFrozen) +{-# NOINLINE symmetricEigenPPar #-} + +-- | Divide-and-conquer specialised symmetric eigenvalue decomposition for @P Double@. +-- Uses raw-primop tridiagonalisation then D&C eigensolver (GEMM-based merge) +-- instead of QR iteration. Faster than 'symmetricEigenP' at larger sizes (n ≥ 50). +symmetricEigenPDC :: forall n. KnownNat n + => Matrix n n M.P Double -> Double + -> (Vector n M.P Double, Matrix n n M.P Double) +symmetricEigenPDC a tol = + let nn = dimVal @n + (q0, diag_, subdiag) = tridiagonalizeP a + (dArr, qArr) = M.withMArrayST (unMatrix q0) $ \mq -> do + md <- M.thawS (unVector diag_) + msd <- M.thawS (unVector subdiag) + let !mbaD = unwrapMutableByteArray md + !offD = unwrapMutableByteArrayOffset md + !mbaSD = unwrapMutableByteArray msd + !offSD = unwrapMutableByteArrayOffset msd + !mbaQ = unwrapMutableByteArray mq + !offQ = unwrapMutableByteArrayOffset mq + dcEigenTridiagOpt mbaD offD mbaSD offSD mbaQ offQ nn 0 (nn - 1) tol + dFrozen <- M.freezeS md + pure (MkVector dFrozen) + in (dArr, MkMatrix qArr) +{-# NOINLINE symmetricEigenPDC #-} + +-- | Parallel QR loop: forks independent sub-problems when a split is found. +-- Operates in IO to enable forkIO for non-overlapping sub-problem ranges. +rawTridiagQRLoopPar :: MutableByteArray RealWorld -> Int + -> MutableByteArray RealWorld -> Int + -> MutableByteArray RealWorld -> Int + -> Int -> Int -> Double -> IO () +rawTridiagQRLoopPar mbaD offD mbaSD offSD mbaQ offQ nn maxIter tol = go 0 0 (nn - 1) + where + rd mba off i = stToIO (readRawD mba off i) + wr mba off i v = stToIO (writeRawD mba off i v) + + go !iter !lo !hi + | iter >= maxIter = pure () + | lo >= hi = pure () + | otherwise = do + sdhi <- rd mbaSD offSD (hi - 1) + dhi1 <- rd mbaD offD (hi - 1) + dhi <- rd mbaD offD hi + if abs sdhi <= tol * (abs dhi1 + abs dhi) + then do wr mbaSD offSD (hi - 1) 0; go iter lo (hi - 1) + else do + sdlo <- rd mbaSD offSD lo + dlo <- rd mbaD offD lo + dlo1 <- rd mbaD offD (lo + 1) + if abs sdlo <= tol * (abs dlo + abs dlo1) + then do wr mbaSD offSD lo 0; go iter (lo + 1) hi + else do + split <- stToIO $ rawFindSplit mbaD offD mbaSD offSD lo hi tol + case split of + Just q -> do + wr mbaSD offSD q 0 + done <- newEmptyMVar + _ <- forkIO $ do + go iter lo q + putMVar done () + go iter (q + 1) hi + takeMVar done + Nothing -> do + let sp1 = sdhi + delta = (dhi1 - dhi) / 2 + sgn = if delta >= 0 then 1 else -1 + shift = dhi - sp1*sp1 / (delta + sgn * sqrt (delta*delta + sp1*sp1)) + stToIO $ rawImplicitQRStep mbaD offD mbaSD offSD mbaQ offQ nn shift lo hi + go (iter + 1) lo hi + +-- | Parallel QR loop with column-major Q for SIMD Givens. +rawTridiagQRLoopParCM :: MutableByteArray RealWorld -> Int + -> MutableByteArray RealWorld -> Int + -> MutableByteArray RealWorld -> Int + -> Int -> Int -> Double -> IO () +rawTridiagQRLoopParCM mbaD offD mbaSD offSD mbaQ offQ nn maxIter tol = go 0 0 (nn - 1) + where + rd mba off i = stToIO (readRawD mba off i) + wr mba off i v = stToIO (writeRawD mba off i v) + + go !iter !lo !hi + | iter >= maxIter = pure () + | lo >= hi = pure () + | otherwise = do + sdhi <- rd mbaSD offSD (hi - 1) + dhi1 <- rd mbaD offD (hi - 1) + dhi <- rd mbaD offD hi + if abs sdhi <= tol * (abs dhi1 + abs dhi) + then do wr mbaSD offSD (hi - 1) 0; go iter lo (hi - 1) + else do + sdlo <- rd mbaSD offSD lo + dlo <- rd mbaD offD lo + dlo1 <- rd mbaD offD (lo + 1) + if abs sdlo <= tol * (abs dlo + abs dlo1) + then do wr mbaSD offSD lo 0; go iter (lo + 1) hi + else do + split <- stToIO $ rawFindSplit mbaD offD mbaSD offSD lo hi tol + case split of + Just q -> do + wr mbaSD offSD q 0 + done <- newEmptyMVar + _ <- forkIO $ do + go iter lo q + putMVar done () + go iter (q + 1) hi + takeMVar done + Nothing -> do + let sp1 = sdhi + delta = (dhi1 - dhi) / 2 + sgn = if delta >= 0 then 1 else -1 + shift = dhi - sp1*sp1 / (delta + sgn * sqrt (delta*delta + sp1*sp1)) + stToIO $ rawImplicitQRStepCM mbaD offD mbaSD offSD mbaQ offQ nn shift lo hi + go (iter + 1) lo hi + +-- | Read a Double from a raw MutableByteArray at element index. +readRawD :: MutableByteArray s -> Int -> Int -> ST s Double +readRawD (MutableByteArray mba) (I# off) (I# i) = ST $ \s -> + case readDoubleArray# mba (off +# i) s of + (# s', v #) -> (# s', D# v #) +{-# INLINE readRawD #-} + +-- | Write a Double to a raw MutableByteArray at element index. +writeRawD :: MutableByteArray s -> Int -> Int -> Double -> ST s () +writeRawD (MutableByteArray mba) (I# off) (I# i) (D# v) = ST $ \s -> + case writeDoubleArray# mba (off +# i) v s of + s' -> (# s', () #) +{-# INLINE writeRawD #-} + +-- | Read an Int from a raw MutableByteArray at element index. +readRawI :: MutableByteArray s -> Int -> Int -> ST s Int +readRawI (MutableByteArray mba) (I# off) (I# i) = ST $ \s -> + case readIntArray# mba (off +# i) s of + (# s', v #) -> (# s', I# v #) +{-# INLINE readRawI #-} + +-- | Write an Int to a raw MutableByteArray at element index. +writeRawI :: MutableByteArray s -> Int -> Int -> Int -> ST s () +writeRawI (MutableByteArray mba) (I# off) (I# i) (I# v) = ST $ \s -> + case writeIntArray# mba (off +# i) v s of + s' -> (# s', () #) +{-# INLINE writeRawI #-} + +-- | Read a Double from an immutable ByteArray at element index. +indexRawD :: ByteArray -> Int -> Int -> Double +indexRawD (ByteArray ba) (I# off) (I# i) = + case indexDoubleArray# ba (off +# i) of + v -> D# v +{-# INLINE indexRawD #-} + +-- | Raw primop QR loop: all diagonal/subdiagonal access via raw ByteArray# primops. +rawTridiagQRLoop :: MutableByteArray s -> Int -- ^ diagonal array + offset + -> MutableByteArray s -> Int -- ^ subdiagonal array + offset + -> MutableByteArray s -> Int -- ^ Q matrix + offset + -> Int -> Int -> Double -> ST s () +rawTridiagQRLoop mbaD offD mbaSD offSD mbaQ offQ nn maxIter tol = go 0 0 (nn - 1) (nn - 1) (0 :: Int) + where + go !iter !lo !hi !lastHi !stall + | iter >= maxIter = pure () + | stall >= 20 = pure () -- bail if 20 consecutive steps fail to deflate + | lo >= hi = pure () + | otherwise = do + -- Bottom deflation + sdhi <- readRawD mbaSD offSD (hi - 1) + dhi1 <- readRawD mbaD offD (hi - 1) + dhi <- readRawD mbaD offD hi + if abs sdhi <= tol * (abs dhi1 + abs dhi) + then do writeRawD mbaSD offSD (hi - 1) 0; go iter lo (hi - 1) hi 0 + else do + -- Top deflation + sdlo <- readRawD mbaSD offSD lo + dlo <- readRawD mbaD offD lo + dlo1 <- readRawD mbaD offD (lo + 1) + if abs sdlo <= tol * (abs dlo + abs dlo1) + then do writeRawD mbaSD offSD lo 0; go iter (lo + 1) hi hi 0 + else do + -- Interior deflation: find split point + split <- rawFindSplit mbaD offD mbaSD offSD lo hi tol + case split of + Just q -> do + writeRawD mbaSD offSD q 0 + go iter lo q hi 0 + go iter (q + 1) hi hi 0 + Nothing -> do + -- AED: scan bottom window before committing to a QR step + newHi <- if hi - lo >= 6 + then rawAEDScan mbaD offD mbaSD offSD tol lo hi + else pure hi + if newHi < hi + then go iter lo newHi hi 0 -- deflated without QR sweep + else do + -- Compute Wilkinson shift from bottom 2×2 block + let !sp1 = sdhi + !delta = (dhi1 - dhi) / 2 + !sgn = if delta >= 0 then 1 else -1 + !shift1 = dhi - sp1*sp1 / (delta + sgn * sqrt (delta*delta + sp1*sp1)) + rawImplicitQRStep mbaD offD mbaSD offSD mbaQ offQ nn shift1 lo hi + -- Double-shift: apply second shift if active range is large enough + when (hi - lo >= 4) $ do + sdhi' <- readRawD mbaSD offSD (hi - 1) + dhi1' <- readRawD mbaD offD (hi - 1) + dhi' <- readRawD mbaD offD hi + when (abs sdhi' > tol * (abs dhi1' + abs dhi')) $ do + let !delta' = (dhi1' - dhi') / 2 + !sgn' = if delta' >= 0 then 1 else -1 + !shift2 = dhi' - sdhi'*sdhi' / (delta' + sgn' * sqrt (delta'*delta' + sdhi'*sdhi')) + rawImplicitQRStep mbaD offD mbaSD offSD mbaQ offQ nn shift2 lo hi + let !newStall = if hi == lastHi then stall + 1 else 0 + go (iter + 1) lo hi hi newStall + +-- | Aggressive Early Deflation: scan bottom w entries for negligible subdiagonals. +-- Returns the new (possibly lower) hi. Deflates from the bottom up, setting +-- negligible subdiagonal entries to zero. +rawAEDScan :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> Double -> Int -> Int -> ST s Int +rawAEDScan mbaD offD mbaSD offSD tol lo hi = scan hi + where + !w = min 6 ((hi - lo + 1) `div` 3) + !bottom = max (lo + 1) (hi - w) + scan !h + | h <= bottom = pure h + | otherwise = do + sdk <- readRawD mbaSD offSD (h - 1) + dk1 <- readRawD mbaD offD (h - 1) + dk <- readRawD mbaD offD h + let !absdk1 = abs dk1 + !absdk = abs dk + !threshold = tol * (absdk1 + absdk) + if abs sdk <= threshold + then do writeRawD mbaSD offSD (h - 1) 0; scan (h - 1) + else pure h +{-# INLINE rawAEDScan #-} + +-- | Raw primop interior split search. +rawFindSplit :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> Int -> Int -> Double -> ST s (Maybe Int) +rawFindSplit mbaD offD mbaSD offSD lo hi tol = scan (lo + 1) + where + scan q + | q >= hi - 1 = pure Nothing + | otherwise = do + sdq <- readRawD mbaSD offSD q + dq <- readRawD mbaD offD q + dq1 <- readRawD mbaD offD (q + 1) + if abs sdq <= tol * (abs dq + abs dq1) + then pure (Just q) + else scan (q + 1) + +-- | Raw primop implicit QR step via bulge-chasing Givens rotations. +rawImplicitQRStep :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> MutableByteArray s -> Int + -> Int -> Double -> Int -> Int -> ST s () +rawImplicitQRStep mbaD offD mbaSD offSD mbaQ offQ nn shift lo hi = do + dlo <- readRawD mbaD offD lo + sdlo <- readRawD mbaSD offSD lo + chase lo (dlo - shift) sdlo + where + chase k x z = do + let (c, s) = givensRotation x z + when (k > lo) $ + writeRawD mbaSD offSD (k-1) (c * x - s * z) + dk <- readRawD mbaD offD k + ek <- readRawD mbaSD offSD k + dk1 <- readRawD mbaD offD (k+1) + writeRawD mbaD offD k (c*c*dk - 2*c*s*ek + s*s*dk1) + writeRawD mbaD offD (k+1) (s*s*dk + 2*c*s*ek + c*c*dk1) + writeRawD mbaSD offSD k (c*s*(dk - dk1) + (c*c - s*s)*ek) + -- Raw primop Givens rotation on Q + rawMutApplyGivensColumns mbaQ offQ nn c (negate s) k (k+1) nn + if k + 1 < hi + then do + ek1 <- readRawD mbaSD offSD (k+1) + let z' = -s * ek1 + writeRawD mbaSD offSD (k+1) (c * ek1) + ek_new <- readRawD mbaSD offSD k + chase (k+1) ek_new z' + else pure () + +-- | Column-major QR loop: same as rawTridiagQRLoop but Q is column-major. +-- In column-major layout, Q[i,j] at off + j*nn + i. This enables SIMD +-- Givens column updates (4 rows at a time via DoubleX4#). +rawTridiagQRLoopCM :: MutableByteArray s -> Int -- ^ diagonal array + offset + -> MutableByteArray s -> Int -- ^ subdiagonal array + offset + -> MutableByteArray s -> Int -- ^ Q matrix (COLUMN-MAJOR) + offset + -> Int -> Int -> Double -> ST s () +rawTridiagQRLoopCM mbaD offD mbaSD offSD mbaQ offQ nn maxIter tol = go 0 0 (nn - 1) (nn - 1) (0 :: Int) + where + go !iter !lo !hi !lastHi !stall + | iter >= maxIter = pure () + | stall >= 20 = pure () -- bail if 20 consecutive steps fail to deflate + | lo >= hi = pure () + | otherwise = do + sdhi <- readRawD mbaSD offSD (hi - 1) + dhi1 <- readRawD mbaD offD (hi - 1) + dhi <- readRawD mbaD offD hi + if abs sdhi <= tol * (abs dhi1 + abs dhi) + then do writeRawD mbaSD offSD (hi - 1) 0; go iter lo (hi - 1) hi 0 + else do + sdlo <- readRawD mbaSD offSD lo + dlo <- readRawD mbaD offD lo + dlo1 <- readRawD mbaD offD (lo + 1) + if abs sdlo <= tol * (abs dlo + abs dlo1) + then do writeRawD mbaSD offSD lo 0; go iter (lo + 1) hi hi 0 + else do + split <- rawFindSplit mbaD offD mbaSD offSD lo hi tol + case split of + Just q -> do + writeRawD mbaSD offSD q 0 + go iter lo q hi 0 + go iter (q + 1) hi hi 0 + Nothing -> do + -- AED: scan bottom window before committing to a QR step + newHi <- if hi - lo >= 6 + then rawAEDScan mbaD offD mbaSD offSD tol lo hi + else pure hi + if newHi < hi + then go iter lo newHi hi 0 -- deflated without QR sweep + else do + -- Compute both eigenvalues of bottom 2×2 block + let !sp1 = sdhi + !delta = (dhi1 - dhi) / 2 + !sgn = if delta >= 0 then 1 else -1 + !shift1 = dhi - sp1*sp1 / (delta + sgn * sqrt (delta*delta + sp1*sp1)) + rawImplicitQRStepCM mbaD offD mbaSD offSD mbaQ offQ nn shift1 lo hi + -- Double-shift: apply second shift if active range is large enough + when (hi - lo >= 4) $ do + -- Check if first shift already deflated the bottom + sdhi' <- readRawD mbaSD offSD (hi - 1) + dhi1' <- readRawD mbaD offD (hi - 1) + dhi' <- readRawD mbaD offD hi + when (abs sdhi' > tol * (abs dhi1' + abs dhi')) $ do + -- Compute the other eigenvalue of the (updated) bottom 2×2 + let !delta' = (dhi1' - dhi') / 2 + !sgn' = if delta' >= 0 then 1 else -1 + !shift2 = dhi' - sdhi'*sdhi' / (delta' + sgn' * sqrt (delta'*delta' + sdhi'*sdhi')) + rawImplicitQRStepCM mbaD offD mbaSD offSD mbaQ offQ nn shift2 lo hi + let !newStall = if hi == lastHi then stall + 1 else 0 + go (iter + 1) lo hi hi newStall + +-- | Column-major implicit QR step: same as rawImplicitQRStep but uses +-- SIMD Givens on column-major Q layout. +rawImplicitQRStepCM :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> MutableByteArray s -> Int + -> Int -> Double -> Int -> Int -> ST s () +rawImplicitQRStepCM mbaD offD mbaSD offSD mbaQ offQ nn shift lo hi = do + dlo <- readRawD mbaD offD lo + sdlo <- readRawD mbaSD offSD lo + chase lo (dlo - shift) sdlo + where + chase k x z = do + let (c, s) = givensRotation x z + when (k > lo) $ + writeRawD mbaSD offSD (k-1) (c * x - s * z) + dk <- readRawD mbaD offD k + ek <- readRawD mbaSD offSD k + dk1 <- readRawD mbaD offD (k+1) + writeRawD mbaD offD k (c*c*dk - 2*c*s*ek + s*s*dk1) + writeRawD mbaD offD (k+1) (s*s*dk + 2*c*s*ek + c*c*dk1) + writeRawD mbaSD offSD k (c*s*(dk - dk1) + (c*c - s*s)*ek) + -- SIMD Givens rotation on column-major Q + rawMutApplyGivensColumnsCM mbaQ offQ nn c (negate s) k (k+1) nn + if k + 1 < hi + then do + ek1 <- readRawD mbaSD offSD (k+1) + let z' = -s * ek1 + writeRawD mbaSD offSD (k+1) (c * ek1) + ek_new <- readRawD mbaSD offSD k + chase (k+1) ek_new z' + else pure () + +-- -------------------------------------------------------------------------- +-- Divide-and-conquer tridiagonal eigensolver (GVL4 Section 8.4) +-- Optimised: pre-allocated workspace, QR fallback for small subproblems, +-- unsafeFreezeByteArray to avoid O(n²) copies. +-- -------------------------------------------------------------------------- + +-- | Optimised divide-and-conquer eigensolver for a symmetric tridiagonal matrix. +-- Pre-allocates all temporary arrays once (eliminating per-level GC pressure), +-- falls back to QR for subproblems ≤ 25, and uses unsafeFreezeByteArray for +-- O(1) GEMM input preparation. +-- +-- The algorithm maintains a LOCAL eigenvector matrix (maxN × maxN, starting as +-- identity) throughout the D&C recursion. The z-vector for each merge step is +-- extracted from this local matrix (not the global Q), ensuring correctness. +-- At the end, the global Q is updated via a single GEMM: Q_out = Q_in * Q_local. +dcEigenTridiagOpt :: MutableByteArray s -> Int -- ^ d + offset + -> MutableByteArray s -> Int -- ^ e + offset + -> MutableByteArray s -> Int -- ^ Q + offset (fullN × fullN row-major) + -> Int -- ^ fullN (Q dimension) + -> Int -> Int -- ^ lo, hi (active range, inclusive) + -> Double -- ^ tolerance + -> ST s () +dcEigenTridiagOpt mbaD offD mbaE offE mbaQ offQ fullN lo0 hi0 tol + | hi0 <= lo0 = pure () + | otherwise = do + let !maxN = hi0 - lo0 + 1 + -- Pre-allocate all workspace at maximum needed size (once, not per-level) + wsLam <- newByteArray (maxN * 8) + wsZ <- newByteArray (maxN * 8) + wsDSort <- newByteArray (maxN * 8) + wsZSort <- newByteArray (maxN * 8) + wsIdx <- newByteArray (maxN * 8) + wsPerm <- newByteArray (maxN * 8) -- deflation permutation (Int indices) + wsW <- newByteArray (maxN * maxN * 8) + wsQsub <- newByteArray (maxN * maxN * 8) + wsResult <- newByteArray (maxN * maxN * 8) + wsQtemp <- newByteArray (maxN * maxN * 8) + + -- Local eigenvector accumulator (maxN × maxN), initialised to identity. + -- Indexed with LOCAL coordinates: row/col in [0..maxN-1]. + -- Local index i corresponds to global index (lo0 + i). + wsQlocal <- newByteArray (maxN * maxN * 8) + rawZeroDoubles wsQlocal 0 (maxN * maxN) + forM_ [0..maxN-1] $ \i -> writeRawD wsQlocal 0 (i * maxN + i) 1 + + let !dcThreshold = 25 + + -- Convenience: convert global index to local index + toLocal !g = g - lo0 + + -- Main recursive function (captures workspace via closure) + -- All operations affect wsQlocal (local eigenvector accumulator), + -- NOT the global Q matrix. + dcGo !lo !hi + | lo >= hi = pure () + | hi == lo + 1 = do + -- 2×2 direct eigensolve + d0 <- readRawD mbaD offD lo + d1 <- readRawD mbaD offD hi + e0 <- readRawD mbaE offE lo + let !tr = d0 + d1 + !det_ = d0 * d1 - e0 * e0 + !disc = sqrt (max 0 (tr * tr - 4 * det_)) + !lam1 = (tr - disc) / 2 + !lam2 = (tr + disc) / 2 + (!c, !s) = if abs e0 < tol * (abs d0 + abs d1) + then (1, 0) + else let !theta = (d1 - d0) / (2 * e0) + !t_ = if theta >= 0 + then 1 / (theta + sqrt (1 + theta * theta)) + else 1 / (theta - sqrt (1 + theta * theta)) + !c_ = 1 / sqrt (1 + t_ * t_) + in (c_, t_ * c_) + writeRawD mbaD offD lo lam1 + writeRawD mbaD offD hi lam2 + writeRawD mbaE offE lo 0 + -- Apply Givens to LOCAL eigenvector matrix columns + rawMutApplyGivensColumns wsQlocal 0 maxN c (negate s) (toLocal lo) (toLocal hi) maxN + | hi - lo + 1 <= dcThreshold = do + -- QR fallback for small subproblems + let !k = hi - lo + 1 + -- Initialise wsQtemp as k×k identity + rawZeroDoubles wsQtemp 0 (k * k) + forM_ [0..k-1] $ \i -> writeRawD wsQtemp 0 (i * k + i) 1 + -- Run QR iteration on d[lo..hi], e[lo..hi-1] + rawTridiagQRLoop mbaD (offD + lo) mbaE (offE + lo) wsQtemp 0 k (30 * k) tol + -- Apply rotation to LOCAL eigenvector matrix: + -- wsQlocal[:, toLocal(lo)..toLocal(lo)+k-1] *= wsQtemp + applyRotToQlocal (toLocal lo) k wsQtemp + | otherwise = do + -- D&C merge for larger subproblems + let !k = (lo + hi) `div` 2 + !n1 = k - lo + 1 + !n2 = hi - k + !nn = hi - lo + 1 + -- Local coordinates for the split + !kL = toLocal k + !loL = toLocal lo + + -- Read and modify the coupling element + beta <- readRawD mbaE offE k + dk <- readRawD mbaD offD k + dk1 <- readRawD mbaD offD (k + 1) + let !absBeta = abs beta + !rho = absBeta + writeRawD mbaD offD k (dk - absBeta) + writeRawD mbaD offD (k+1) (dk1 - absBeta) + writeRawD mbaE offE k 0 + + -- Recurse on T1 [lo..k] and T2 [k+1..hi] + dcGo lo k + dcGo (k + 1) hi + + -- === Merge phase === + -- Extract z vector from LOCAL eigenvector matrix rows. + -- z[0..n1-1] = last row of Q1 = row kL of wsQlocal, columns loL..loL+n1-1 + -- z[n1..nn-1] = first row of Q2 = row (kL+1) of wsQlocal, columns loL+n1..loL+nn-1 + forM_ [0..n1-1] $ \i -> do + qv <- readRawD wsQlocal 0 (kL * maxN + (loL + i)) + writeRawD wsZ 0 i qv + forM_ [0..n2-1] $ \i -> do + qv <- readRawD wsQlocal 0 ((kL + 1) * maxN + (loL + n1 + i)) + let !zv = if beta < 0 then negate qv else qv + writeRawD wsZ 0 (n1 + i) zv + + -- Copy d[lo..hi] and z into sortable arrays with indices + forM_ [0..nn-1] $ \i -> do + di <- readRawD mbaD offD (lo + i) + writeRawD wsDSort 0 i di + writeRawD wsZSort 0 i =<< readRawD wsZ 0 i + writeRawD wsIdx 0 i (fromIntegral i) + + -- Sort by d values (insertion sort) + forM_ [1..nn-1] $ \i -> do + di <- readRawD wsDSort 0 i + zi <- readRawD wsZSort 0 i + idxi <- readRawD wsIdx 0 i + let insertAt !j + | j < 0 = do + writeRawD wsDSort 0 0 di + writeRawD wsZSort 0 0 zi + writeRawD wsIdx 0 0 idxi + | otherwise = do + dj <- readRawD wsDSort 0 j + if dj > di + then do + writeRawD wsDSort 0 (j+1) dj + writeRawD wsZSort 0 (j+1) =<< readRawD wsZSort 0 j + writeRawD wsIdx 0 (j+1) =<< readRawD wsIdx 0 j + insertAt (j - 1) + else do + writeRawD wsDSort 0 (j+1) di + writeRawD wsZSort 0 (j+1) zi + writeRawD wsIdx 0 (j+1) idxi + insertAt (i - 1) + + -- Close-d deflation (cf. LAPACK dlaed2): when consecutive + -- sorted d values are nearly equal, a Givens rotation zeros + -- one z entry, preventing ill-conditioned secular roots. + dMaxAbs <- readRawD wsDSort 0 (nn - 1) + dMinAbs <- readRawD wsDSort 0 0 + let !closeDTol = 8 * 2.220446049250313e-16 + * max (abs dMaxAbs) (abs dMinAbs + rho) + forM_ [0..nn-2] $ \i -> do + di <- readRawD wsDSort 0 i + di1 <- readRawD wsDSort 0 (i + 1) + when (abs (di1 - di) <= closeDTol) $ do + zi <- readRawD wsZSort 0 i + zi1 <- readRawD wsZSort 0 (i + 1) + let !r = sqrt (zi * zi + zi1 * zi1) + when (r > 1e-300) $ do + let !c = zi1 / r + !s = zi / r + -- Zero z[i], combine into z[i+1] + writeRawD wsZSort 0 i 0 + writeRawD wsZSort 0 (i + 1) r + -- Apply same Givens to eigenvector columns + origI <- readRawD wsIdx 0 i + origI1 <- readRawD wsIdx 0 (i + 1) + let !colI = loL + (round origI :: Int) + !colI1 = loL + (round origI1 :: Int) + rawMutApplyGivensColumns wsQlocal 0 maxN c s colI colI1 maxN + + -- Deflation: partition into non-deflated and deflated. + -- The perturbation-based criterion ensures that entries whose + -- eigenvalue shift rho*z[i]² is below machine precision + -- relative to |d[i]| are deflated, preventing the eigenvector + -- formula from dividing by zero (d[j] - lambda == 0 in FP). + zn2 <- sumZSq wsZSort 0 nn + let !eps_ = 2.220446049250313e-16 + !matNorm = max (abs dMaxAbs) (abs dMinAbs) + rho * zn2 + !basicDeflTol = max (tol * sqrt zn2) (8 * eps_ * matNorm) + -- Perturbation deflation: deflate when rho*z² do + origIdx <- readRawD wsIdx 0 sortedJ + let !origJ = round origIdx :: Int + !srcCol = loL + origJ + rawCopyColumn wsQlocal 0 maxN srcCol wsQsub 0 nn sortedJ maxN + + if kND == 0 + then do + -- All deflated: eigenvalues = sorted d, eigenvectors = sorted Qlocal cols + forM_ [0..nn-1] $ \i -> + rawCopyColumn wsQsub 0 nn i wsQlocal 0 maxN (loL + i) maxN + forM_ [0..nn-1] $ \i -> do + di <- readRawD wsDSort 0 i + writeRawD mbaD offD (lo + i) di + + else if kND == nn + then do + -- No deflation: full secular solve + full GEMM + secularSolve wsLam 0 wsDSort 0 wsZSort 0 rho nn deflTol + dcEigenvectors wsW 0 wsDSort 0 wsZSort 0 wsLam 0 rho nn + + baQsub <- unsafeFreezeByteArray wsQsub + baW <- unsafeFreezeByteArray wsW + rawZeroDoubles wsResult 0 (maxN * nn) + rawGemmKernel baQsub 0 baW 0 wsResult 0 maxN nn nn + + forM_ [0..nn-1] $ \i -> + rawCopyColumn wsResult 0 nn i wsQlocal 0 maxN (loL + i) maxN + forM_ [0..nn-1] $ \i -> do + lam <- readRawD wsLam 0 i + writeRawD mbaD offD (lo + i) lam + + else do + -- Partial deflation: reduced secular solve + reduced GEMM + -- Build compressed d_nd[0..kND-1] and z_nd[0..kND-1] + -- Store in wsQtemp: d_nd at offset 0, z_nd at offset kND + forM_ [0..kND-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + dpi <- readRawD wsDSort 0 pi_ + zpi <- readRawD wsZSort 0 pi_ + writeRawD wsQtemp 0 j dpi + writeRawD wsQtemp 0 (kND + j) zpi + + -- Solve kND secular equations on compressed system + secularSolve wsLam 0 wsQtemp 0 wsQtemp kND rho kND deflTol + + -- Compute kND×kND eigenvector matrix W_nd + dcEigenvectors wsW 0 wsQtemp 0 wsQtemp kND wsLam 0 rho kND + + -- Copy deflated columns from wsQsub to wsQlocal + forM_ [kND..nn-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + rawCopyColumn wsQsub 0 nn pi_ wsQlocal 0 maxN (loL + j) maxN + + -- Extract Q_nd (maxN×kND) from non-deflated columns + forM_ [0..kND-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + rawCopyColumn wsQsub 0 nn pi_ wsResult 0 kND j maxN + + -- GEMM: wsQsub(maxN×kND) = Q_nd(maxN×kND) × W_nd(kND×kND) + baQnd <- unsafeFreezeByteArray wsResult + baW <- unsafeFreezeByteArray wsW + rawZeroDoubles wsQsub 0 (maxN * kND) + rawGemmKernel baQnd 0 baW 0 wsQsub 0 maxN kND kND + + -- Copy GEMM result (non-deflated columns) to wsQlocal + forM_ [0..kND-1] $ \j -> + rawCopyColumn wsQsub 0 kND j wsQlocal 0 maxN (loL + j) maxN + + -- Write eigenvalues: non-deflated from wsLam, deflated from wsDSort + forM_ [0..kND-1] $ \i -> do + lam <- readRawD wsLam 0 i + writeRawD mbaD offD (lo + i) lam + forM_ [kND..nn-1] $ \j -> do + pi_ <- readRawI wsPerm 0 j + di <- readRawD wsDSort 0 pi_ + writeRawD mbaD offD (lo + j) di + + -- Apply a k×k rotation matrix to wsQlocal columns [colOff..colOff+k-1] via GEMM + -- colOff is in LOCAL coordinates. + applyRotToQlocal !colOff !k rotMat = do + -- Extract wsQlocal[:, colOff..colOff+k-1] into wsQsub (maxN × k) + forM_ [0..k-1] $ \j -> + rawCopyColumn wsQlocal 0 maxN (colOff + j) wsQsub 0 k j maxN + -- O(1) freeze for GEMM inputs + baQsub <- unsafeFreezeByteArray wsQsub + baRot <- unsafeFreezeByteArray rotMat + -- Zero result + rawZeroDoubles wsResult 0 (maxN * k) + -- GEMM: result(maxN×k) = Qsub(maxN×k) * Rot(k×k) + rawGemmKernel baQsub 0 baRot 0 wsResult 0 maxN k k + -- Copy result back to wsQlocal + forM_ [0..k-1] $ \j -> + rawCopyColumn wsResult 0 k j wsQlocal 0 maxN (colOff + j) maxN + + -- Run the D&C recursion (operates on wsQlocal and mbaD/mbaE) + dcGo lo0 hi0 + + -- Final step: apply wsQlocal to global Q via GEMM + -- Q[:, lo0..hi0] = Q[:, lo0..hi0] * wsQlocal + forM_ [0..maxN-1] $ \j -> + rawCopyColumn mbaQ offQ fullN (lo0 + j) wsQsub 0 maxN j fullN + + baQsub <- unsafeFreezeByteArray wsQsub + baQlocal <- unsafeFreezeByteArray wsQlocal + rawZeroDoubles wsResult 0 (fullN * maxN) + rawGemmKernel baQsub 0 baQlocal 0 wsResult 0 fullN maxN maxN + + forM_ [0..maxN-1] $ \j -> + rawCopyColumn wsResult 0 maxN j mbaQ offQ fullN (lo0 + j) fullN + +-- | Solve the secular equation: f(λ) = 1 + ρ * Σ z[i]² / (d[i] - λ) = 0 +-- for all nn roots. Roots are stored in mbaLam. +-- d must be sorted in ascending order. +secularSolve :: MutableByteArray s -> Int -- ^ output eigenvalues + -> MutableByteArray s -> Int -- ^ sorted d (poles) + -> MutableByteArray s -> Int -- ^ sorted z + -> Double -> Int -> Double -- ^ rho, n, deflation tolerance + -> ST s () +secularSolve mbaLam offLam mbaD offD mbaZ offZ rho nn deflTol = do + forM_ [0..nn-1] $ \i -> do + zi <- readRawD mbaZ offZ i + di <- readRawD mbaD offD i + if abs zi <= deflTol + then do + -- Deflated: eigenvalue = d[i] + writeRawD mbaLam offLam i di + else do + -- For small z[i], use first-order perturbation formula directly. + -- This avoids the iterative solver's difficulty with nearly-flat + -- secular functions near d[i]. + let !zi2 = zi * zi + !pertTol = sqrt (2.220446049250313e-16) * (1 + abs di) + if abs zi < pertTol + then do + -- Perturbation: lambda ≈ d[i] + rho * z[i]² / (1 + rho * Σ_{j≠i} z[j]²/(d[j]-d[i])) + farSum <- farPoleSumSkip mbaD offD mbaZ offZ nn i di + let !denom = 1 + rho * farSum + !delta = rho * zi2 / denom + writeRawD mbaLam offLam i (di + delta) + else do + lam <- secularSolveOne mbaD offD mbaZ offZ rho i nn + writeRawD mbaLam offLam i lam + +-- | Solve one root of the secular equation between d[j] and d[j+1] +-- (or between d[n-1] and +infinity for the last root when rho > 0, +-- or between -infinity and d[0] for the first root when rho < 0). +-- Uses the Gragg/Borges fixed-weight quadratic method (cf. LAPACK dlasd4): +-- splits f(λ) at the two closest poles, approximates far terms as constant, +-- and solves the resulting quadratic for rapid convergence (2–4 iterations). +secularSolveOne :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> Double -> Int -> Int -> ST s Double +secularSolveOne mbaD offD mbaZ offZ rho j nn = do + dj <- readRawD mbaD offD j + zj <- readRawD mbaZ offZ j + let !zj2 = zj * zj + -- Determine bracket and second pole + if rho > 0 + then if j < nn - 1 + then do + dj1 <- readRawD mbaD offD (j+1) + zj1 <- readRawD mbaZ offZ (j+1) + -- Interior root between d[j] and d[j+1] + let !gap = dj1 - dj + !mid = dj + gap * 0.5 + lam0 <- fixedWeightLoop 0 mid dj dj1 dj dj1 gap zj2 (zj1 * zj1) + -- Polish with Newton iterations for higher accuracy + newtonPolish 0 lam0 dj dj1 + else do + -- Last root when rho > 0: between d[n-1] and d[n-1] + rho*||z||² + zn2 <- sumZSq mbaZ offZ nn + -- Compute better initial guess via perturbation theory: + -- f(d[nn-1]+δ) = 0 ⟹ δ ≈ rho * z[nn-1]² / (1 + rho * Σ_{i 1e-300 + then rho * zj2 / denominator + else rho * zn2 + !hi_ = dj + max (rho * zn2) (2 * delta0) + !mid = dj + max delta0 (1e-15 * (1 + abs dj)) + newtonLoop 0 mid dj hi_ + else if j > 0 + then do + dj0 <- readRawD mbaD offD (j-1) + zj0 <- readRawD mbaZ offZ (j-1) + -- Interior root between d[j-1] and d[j] (rho < 0) + let !gap = dj - dj0 + !mid = dj0 + gap * 0.5 + fixedWeightLoop 0 mid dj0 dj dj0 dj gap (zj0 * zj0) zj2 + else do + -- First root when rho < 0 + zn2 <- sumZSq mbaZ offZ nn + farSum <- farPoleSum mbaD offD mbaZ offZ nn dj + let !denominator = 1 + rho * farSum + !delta0 = if abs denominator > 1e-300 + then abs rho * zj2 / abs denominator + else abs rho * zn2 + !lo_ = dj - max (abs rho * zn2) (2 * delta0) + !mid = dj - max delta0 (1e-15 * (1 + abs dj)) + newtonLoop 0 mid lo_ dj + where + !maxIter_ = 100 :: Int + + -- Newton polishing: 3 Newton steps to refine eigenvalue to machine precision. + -- Uses the secular function and its derivative for rapid convergence. + newtonPolish !iter !lam !lb !ub + | iter >= 3 = pure lam + | otherwise = do + (f, fp) <- secularFuncAndDeriv mbaD offD mbaZ offZ rho nn lam + if abs f < 1e-15 * (1 + abs lam) || abs fp < 1e-300 + then pure lam + else do + let !step = f / fp + !lamNew = lam - step + !clamped = max lb (min ub lamNew) + if abs (clamped - lam) < 1e-16 * (1 + abs lam) + then pure clamped + else newtonPolish (iter + 1) clamped lb ub + + -- Fixed-weight quadratic iteration for interior roots. + -- dLo, dHi are the two FIXED closest poles (never change during iteration). + -- lb, ub are the bracket bounds (narrow during iteration). + -- gap = dHi - dLo. z2Lo, z2Hi are z²[lo_pole] and z²[hi_pole]. + fixedWeightLoop !iter !lam !lb !ub !dLo !dHi !gap !z2Lo !z2Hi + | iter >= maxIter_ = pure lam + | otherwise = do + -- Evaluate f(λ) with split at the fixed poles dLo and dHi + (psiSum, phiSum) <- secularFuncSplit mbaD offD mbaZ offZ nn lam dLo dHi + let !f = 1 + rho * (psiSum + phiSum) + if abs f < 1e-15 * (1 + abs lam) + then pure lam + else do + -- Extract close-pole contributions using FIXED poles + let !deltaLo = dLo - lam -- fixed pole - λ (negative for interior root) + !deltaHi = dHi - lam -- fixed pole - λ (positive for interior root) + -- Protect against division by zero near poles + !aClose = if abs deltaLo > 1e-300 then z2Lo / deltaLo else 0 + !bClose = if abs deltaHi > 1e-300 then z2Hi / deltaHi else 0 + -- "Far" residual: W = f - ρ*(aClose + bClose) + !w = f - rho * (aClose + bClose) + -- Quadratic in τ = dLo - λ (= deltaLo): + -- W*τ² - (W*gap + ρ*z2Lo + ρ*z2Hi)*τ + ρ*z2Lo*gap = 0 + !qa = w + !qb = -(w * gap + rho * z2Lo + rho * z2Hi) + !qc = rho * z2Lo * gap + !disc = qb * qb - 4 * qa * qc + if disc < 0 || abs qa < 1e-300 + then do + -- Degenerate: fall back to bisection + let !(lb', ub') = if f * rho > 0 then (lb, lam) else (lam, ub) + !lamNew = (lb' + ub') * 0.5 + fixedWeightLoop (iter + 1) lamNew lb' ub' dLo dHi gap z2Lo z2Hi + else do + let !sqrtDisc = sqrt disc + -- Two roots for τ = dLo - λ, i.e. λ = dLo - τ + -- Use the numerically stable form + !tauA = if qb <= 0 + then (-qb + sqrtDisc) / (2 * qa) + else 2 * qc / (-qb + sqrtDisc) + !tauB = if qb <= 0 + then 2 * qc / (-qb + sqrtDisc) + else (-qb + sqrtDisc) / (2 * qa) + -- λ = dLo - τ; pick the root in bracket + !lamA = dLo - tauA + !lamB = dLo - tauB + !lamNew0 = if lamA > lb && lamA < ub then lamA + else if lamB > lb && lamB < ub then lamB + else (lb + ub) * 0.5 -- bisection fallback + -- Update bracket + !(lb', ub') = if f * rho > 0 then (lb, lam) else (lam, ub) + -- Ensure lamNew is in updated bracket + !lamNew = if lamNew0 > lb' && lamNew0 < ub' + then lamNew0 + else (lb' + ub') * 0.5 + if abs (lamNew - lam) < 1e-15 * (1 + abs lam) + then pure lamNew + else fixedWeightLoop (iter + 1) lamNew lb' ub' dLo dHi gap z2Lo z2Hi + + -- Newton+bisection fallback for edge roots (first/last eigenvalue). + newtonLoop !iter !lam !lb !ub + | iter >= maxIter_ = pure lam + | otherwise = do + (f, fp) <- secularFuncAndDeriv mbaD offD mbaZ offZ rho nn lam + if abs f < 1e-15 * (1 + abs lam) + then pure lam + else do + let !(lb', ub') = if f > 0 then (lb, lam) else (lam, ub) + !step = f / fp + !lamNew0 = lam - step + !lamNew = if lamNew0 <= lb' || lamNew0 >= ub' + then (lb' + ub') * 0.5 + else lamNew0 + if abs (lamNew - lam) < 1e-15 * (1 + abs lam) + then pure lamNew + else newtonLoop (iter + 1) lamNew lb' ub' + +-- | Evaluate the secular function split at the two bracket poles. +-- Returns (ψ, φ) where f(λ) = 1 + ρ*(ψ + φ). +-- ψ = Σ_{d[i] ≤ dLo} z[i]²/(d[i] - λ), φ = Σ_{d[i] ≥ dHi} z[i]²/(d[i] - λ) +secularFuncSplit :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> Int -> Double -> Double -> Double -> ST s (Double, Double) +secularFuncSplit mbaD offD mbaZ offZ nn lam dLo _dHi = go 0 0 0 + where + go i !psiAcc !phiAcc + | i >= nn = pure (psiAcc, phiAcc) + | otherwise = do + di <- readRawD mbaD offD i + zi <- readRawD mbaZ offZ i + let !diff = di - lam + !zi2 = zi * zi + if abs diff < 1e-300 + then go (i+1) psiAcc phiAcc + else let !term = zi2 / diff + in if di <= dLo + then go (i+1) (psiAcc + term) phiAcc + else go (i+1) psiAcc (phiAcc + term) + +-- | Evaluate the secular function f(λ) = 1 + ρ * Σ z[i]² / (d[i] - λ) +-- and its derivative f'(λ) = ρ * Σ z[i]² / (d[i] - λ)². +secularFuncAndDeriv :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> Double -> Int -> Double -> ST s (Double, Double) +secularFuncAndDeriv mbaD offD mbaZ offZ rho nn lam = do + (fSum, fpSum) <- go 0 0 0 + pure (1 + rho * fSum, rho * fpSum) + where + go i !fAcc !fpAcc + | i >= nn = pure (fAcc, fpAcc) + | otherwise = do + di <- readRawD mbaD offD i + zi <- readRawD mbaZ offZ i + let diff = di - lam + zi2 = zi * zi + if abs diff < 1e-300 + then go (i+1) fAcc fpAcc -- skip near-pole + else go (i+1) (fAcc + zi2 / diff) (fpAcc + zi2 / (diff * diff)) + +-- | Sum of squares of z vector. SIMD-accelerated with DoubleX4# accumulator. +sumZSq :: MutableByteArray s -> Int -> Int -> ST s Double +sumZSq mbaZ offZ nn + | nn < 4 = goScalar 0 0.0 + | otherwise = do + baZ <- unsafeFreezeByteArray mbaZ + let !(ByteArray baZ#) = baZ + !(I# offZ#) = offZ + !nn4 = nn - (nn `rem` 4) + z4 = broadcastDoubleX4# 0.0## + goSimd !i acc4 + | i >= nn4 = do + let !(# a, b, c, d #) = unpackDoubleX4# acc4 + goScalar nn4 (D# a + D# b + D# c + D# d) + | otherwise = + let !(I# ii) = i + zv = indexDoubleArrayAsDoubleX4# baZ# (offZ# +# ii) + !p = timesDoubleX4# zv zv + in goSimd (i + 4) (plusDoubleX4# acc4 p) + goSimd (0 :: Int) z4 + where + goScalar i !acc + | i >= nn = pure acc + | otherwise = do + zi <- readRawD mbaZ offZ i + goScalar (i+1) (acc + zi * zi) + +-- | Sum of z[i]^2 / (d[i] - dj) for i in [0..skip-1], skipping near-zero denominators. +farPoleSum :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> Int -> Double -> ST s Double +farPoleSum mbaD offD mbaZ offZ skip dj = go 0 0 + where + go i !acc + | i >= skip = pure acc + | otherwise = do + di <- readRawD mbaD offD i + zi <- readRawD mbaZ offZ i + let !diff = di - dj + if abs diff < 1e-300 + then go (i+1) acc + else go (i+1) (acc + zi * zi / diff) + +-- | Sum of z[k]^2 / (d[k] - dj) for all k in [0..nn-1] except k == skip. +farPoleSumSkip :: MutableByteArray s -> Int -> MutableByteArray s -> Int + -> Int -> Int -> Double -> ST s Double +farPoleSumSkip mbaD offD mbaZ offZ nn skip dj = go 0 0 + where + go i !acc + | i >= nn = pure acc + | i == skip = go (i+1) acc + | otherwise = do + di <- readRawD mbaD offD i + zi <- readRawD mbaZ offZ i + let !diff = di - dj + if abs diff < 1e-300 + then go (i+1) acc + else go (i+1) (acc + zi * zi / diff) + +-- | Partition sorted indices into non-deflated (|z[i]| > deflTol) and deflated. +-- Returns k (non-deflated count). +-- perm[0..k-1] = sorted indices of non-deflated entries (in sorted order). +-- perm[k..nn-1] = sorted indices of deflated entries (in sorted order). +deflatePartition :: MutableByteArray s -> Int -- ^ sorted z + offset + -> MutableByteArray s -> Int -- ^ output perm (Int array) + offset + -> Int -> Double -- ^ nn, deflTol + -> ST s Int +deflatePartition mbaZ offZ mbaPerm offPerm nn deflTol = do + k <- goND 0 0 + goDF 0 k + pure k + where + goND !i !kND + | i >= nn = pure kND + | otherwise = do + zi <- readRawD mbaZ offZ i + if abs zi > deflTol + then do + writeRawI mbaPerm offPerm kND i + goND (i+1) (kND+1) + else goND (i+1) kND + goDF !i !pos + | i >= nn = pure () + | otherwise = do + zi <- readRawD mbaZ offZ i + if abs zi <= deflTol + then do + writeRawI mbaPerm offPerm pos i + goDF (i+1) (pos+1) + else goDF (i+1) pos + +-- | Compute eigenvector matrix W from secular equation solutions. +-- W[j,i] = z[j] / (d[j] - lambda[i]), each column normalised. +-- Single-pass: writes unnormalised entries and accumulates norm² simultaneously, +-- then normalises each column with SIMD. +-- | Compute eigenvector matrix W for the D&C merge step using the +-- Gu-Eisenstat formula (GVL4 Theorem 8.4.4, p. 469) for improved +-- numerical stability. +-- +-- Instead of the naive W[j,i] = z[j]/(d[j]-λ[i]) which suffers from +-- catastrophic cancellation when d[j] ≈ λ[i], we first compute: +-- +-- z_new[j]² = ∏_k (λ[k] - d[j]) / ∏_{k≠j} (d[k] - d[j]) +-- +-- This is an algebraic identity but computes z_new to full relative accuracy +-- because all factors in numerator and denominator are well-separated. +-- Then W[j,i] = z_new[j] / (d[j] - λ[i]) with column normalization. +dcEigenvectors :: MutableByteArray s -> Int -- ^ W (nn × nn output) + -> MutableByteArray s -> Int -- ^ d (sorted poles) + -> MutableByteArray s -> Int -- ^ z + -> MutableByteArray s -> Int -- ^ lambda (eigenvalues) + -> Double -> Int -- ^ rho, nn + -> ST s () +dcEigenvectors mbaW offW mbaD offD mbaZ offZ mbaLam offLam _rho nn = do + -- Phase 1: Compute z_new via Gu-Eisenstat formula in log space + mbaZnew <- newByteArray (nn * 8) + forM_ [0..nn-1] $ \j -> do + zj <- readRawD mbaZ offZ j + dj <- readRawD mbaD offD j + -- log|z_new[j]²| = Σ_k log|λ[k] - d[j]| - Σ_{k≠j} log|d[k] - d[j]| + logNumer <- goLogSum mbaLam offLam 0 nn dj 0 (-1) -- sum all k + logDenom <- goLogSum mbaD offD 0 nn dj 0 j -- sum all k ≠ j + let !logZ2 = logNumer - logDenom + !absZnew = exp (logZ2 * 0.5) + !znew = if zj >= 0 then absZnew else negate absZnew + writeRawD mbaZnew 0 j znew + + -- Phase 2: Build W[j,i] = z_new[j] / (d[j] - λ[i]), normalise columns + forM_ [0..nn-1] $ \i -> do + lami <- readRawD mbaLam offLam i + norm2 <- writeAndNorm mbaZnew lami i 0 0 + let !invNorm = if norm2 > 0 then 1 / sqrt norm2 else 1 + forM_ [0..nn-1] $ \j -> do + wji <- readRawD mbaW offW (j * nn + i) + writeRawD mbaW offW (j * nn + i) (wji * invNorm) + where + -- Sum of log|arr[k] - val| for k in [lo..hi-1], skipping index 'skip' (-1 = skip none) + goLogSum !arr !off !lo !hi !val !acc !skip + | lo >= hi = pure acc + | lo == skip = goLogSum arr off (lo + 1) hi val acc skip + | otherwise = do + ak <- readRawD arr off lo + let !diff = abs (ak - val) + !logDiff = if diff < 1e-300 then -690.7755 else log diff -- log(1e-300) + goLogSum arr off (lo + 1) hi val (acc + logDiff) skip + + writeAndNorm !mbaZnew !lami !i !j !acc + | j >= nn = pure acc + | otherwise = do + znewj <- readRawD mbaZnew 0 j + dj <- readRawD mbaD offD j + let !diff = dj - lami + !w = if abs diff < 1e-300 then 0 else znewj / diff + writeRawD mbaW offW (j * nn + i) w + writeAndNorm mbaZnew lami i (j + 1) (acc + w * w) + +-- | Classical Jacobi eigenvalue method (GVL4 Section 8.5). +jacobiEigen :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Int -> e -> (Vector n r e, Matrix n n r e) +jacobiEigen a maxSweeps tol = + let nn = dimVal @n + (eigvals, qArr) = M.withMArrayST (unMatrix (identityMatrix @n @r)) $ \mq -> do + ma <- M.thawS (unMatrix a) + jacobiLoop ma mq nn maxSweeps tol + evs <- mapM (\i -> M.readM ma (i :. i)) [0..nn-1] + pure (makeVector @n @r $ \i -> evs !! i) + in (eigvals, MkMatrix qArr) + +jacobiLoop :: (M.Manifest r e, Floating e, Ord e) + => M.MArray s r Ix2 e -> M.MArray s r Ix2 e -> Int -> Int -> e -> ST s () +jacobiLoop ma mq nn maxSweeps tol = go 0 + where + go !sweep + | sweep >= maxSweeps = pure () + | otherwise = do + offNorm <- offDiagNormST ma nn + if offNorm < tol then pure () + else do + forM_ [(p_, q_) | p_ <- [0..nn-2], q_ <- [p_+1..nn-1]] $ \(p_, q_) -> do + apq <- M.readM ma (p_ :. q_) + when (abs apq > tol * 1e-3) $ do + app <- M.readM ma (p_ :. p_) + aqq <- M.readM ma (q_ :. q_) + let (c, s) = jacobiRotation app apq aqq + applyJacobiInPlace ma c s p_ q_ nn + applyGivensRightQ mq c s p_ q_ nn + go (sweep + 1) + +offDiagNormST :: (M.Manifest r e, Floating e) => M.MArray s r Ix2 e -> Int -> ST s e +offDiagNormST ma nn = do + s <- go 0 0 0 + pure (sqrt s) + where go !i !j !acc + | i >= nn = pure acc + | j >= nn = go (i+1) 0 acc + | i == j = go i (j+1) acc + | otherwise = do v <- M.readM ma (i :. j); go i (j+1) (acc + v*v) + +jacobiRotation :: (Floating e, Ord e) => e -> e -> e -> (e, e) +jacobiRotation app apq aqq + | apq == 0 = (1, 0) + | otherwise = + let tau = (aqq - app) / (2 * apq) + t = if tau >= 0 + then 1 / (tau + sqrt (1 + tau * tau)) + else 1 / (tau - sqrt (1 + tau * tau)) + c = 1 / sqrt (1 + t * t) + s = t * c + in (c, s) + +applyJacobiInPlace :: (M.Manifest r e, Num e) + => M.MArray s r Ix2 e -> e -> e -> Int -> Int -> Int -> ST s () +applyJacobiInPlace ma c s p q nn = do + app <- M.readM ma (p :. p) + apq_ <- M.readM ma (p :. q) + aqq <- M.readM ma (q :. q) + M.write_ ma (p :. p) (c*c*app - 2*s*c*apq_ + s*s*aqq) + M.write_ ma (q :. q) (s*s*app + 2*s*c*apq_ + c*c*aqq) + M.write_ ma (p :. q) 0 + M.write_ ma (q :. p) 0 + forM_ [0..nn-1] $ \i -> when (i /= p && i /= q) $ do + aip <- M.readM ma (i :. p) + aiq <- M.readM ma (i :. q) + let aip_new = c * aip - s * aiq + aiq_new = s * aip + c * aiq + M.write_ ma (i :. p) aip_new + M.write_ ma (p :. i) aip_new + M.write_ ma (i :. q) aiq_new + M.write_ ma (q :. i) aiq_new diff --git a/src/Numeric/LinearAlgebra/Massiv/Internal.hs b/src/Numeric/LinearAlgebra/Massiv/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Internal.hs @@ -0,0 +1,261 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Internal +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Internal module providing unsafe constructors, dimension reification helpers, +-- and array creation utilities. This module is re-exported by +-- "Numeric.LinearAlgebra.Massiv" for convenience, but the unsafe constructors +-- ('unsafeMatrix', 'unsafeVector') bypass the dimension checks provided by the +-- smart constructors in "Numeric.LinearAlgebra.Massiv.Types". +-- +-- = Array creation patterns +-- +-- Two primary patterns are provided for constructing dimensioned arrays: +-- +-- 1. __Pure indexed construction__ via 'makeMatrix' and 'makeVector': supply a +-- pure function @(Int -> Int -> e)@ or @(Int -> e)@ that computes each element +-- from its indices. These use massiv's delayed (@D@) intermediate representation +-- and then @'Data.Massiv.Array.compute'@ to materialise the result. +-- +-- 2. __Mutable ST construction__ via 'createMatrix' and 'createVector': supply an +-- @ST@ action operating on a mutable array. This is essential for algorithms +-- that require in-place updates (e.g., LU factorization, Cholesky). +-- +-- Both patterns produce arrays with the sequential ('Data.Massiv.Array.Seq') +-- computation strategy by default. Use the @Comp@ variants ('makeMatrixComp', +-- 'makeVectorComp') for parallel construction. +module Numeric.LinearAlgebra.Massiv.Internal + ( -- * Unsafe constructors + unsafeMatrix + , unsafeVector + -- * Dimension reification + , dimVal + , dimVal2 + , reifyDim + , reifyDim2 + -- * Array creation helpers (pure, sequential) + , makeMatrix + , makeVector + -- * Array creation helpers (pure, with Comp) + , makeMatrixComp + , makeVectorComp + -- * Array creation helpers (mutable ST) + , createMatrix + , createVector + , createMatrixComp + , createVectorComp + -- * Mutable modification helpers + , withMutableMatrix + , withMutableVector + , withMutableMatrix_ + , withMutableVector_ + -- * Indexing helpers + , (!) + , (!.) + -- * Identity and zero + , identityMatrix + , zeroMatrix + , zeroVector + ) where + +import Data.Massiv.Array (Array, Ix2(..), Sz(..), Ix1, Comp(..), D) +import qualified Data.Massiv.Array as M +import GHC.TypeNats (Nat, KnownNat, natVal, SomeNat(..), someNatVal) +import Data.Proxy (Proxy(..)) +import Control.Monad.ST (ST) + +import Numeric.LinearAlgebra.Massiv.Types + +-- | Unsafe matrix constructor — wraps a massiv array with /no/ dimension check. +-- +-- __Precondition__: the array must have exactly \(m\) rows and \(n\) columns. +-- Violating this precondition leads to index-out-of-bounds errors at runtime. +-- +-- Prefer the safe 'matrix' constructor from "Numeric.LinearAlgebra.Massiv.Types" +-- unless you can guarantee correctness (e.g., the array was just constructed +-- with the correct dimensions). +unsafeMatrix :: Array r Ix2 e -> Matrix m n r e +unsafeMatrix = MkMatrix + +-- | Unsafe vector constructor — wraps a massiv array with /no/ size check. +-- +-- __Precondition__: the array must have exactly \(n\) elements. +unsafeVector :: Array r Ix1 e -> Vector n r e +unsafeVector = MkVector + +-- | Get the runtime value of a type-level dimension. +-- +-- @ +-- dimVal \@3 == 3 +-- dimVal \@100 == 100 +-- @ +dimVal :: forall n. KnownNat n => Int +dimVal = fromIntegral (natVal (Proxy @n)) + +-- | Get both dimensions of a matrix type as a tuple. +dimVal2 :: forall m n. (KnownNat m, KnownNat n) => (Int, Int) +dimVal2 = (dimVal @m, dimVal @n) + +-- | Index into a matrix (0-based, unchecked). +-- +-- @mat '!' (i, j)@ returns the element at row \(i\), column \(j\). +-- +-- __Warning__: No bounds checking is performed. Out-of-bounds access +-- results in undefined behaviour for unboxed\/primitive representations. +(!) :: M.Manifest r e => Matrix m n r e -> (Int, Int) -> e +(!) (MkMatrix arr) (i, j) = M.index' arr (i :. j) + +-- | Index into a vector (0-based, unchecked). +-- +-- @vec '!.' i@ returns the element at position \(i\). +(!.) :: M.Manifest r e => Vector n r e -> Int -> e +(!.) (MkVector arr) i = M.index' arr i + +-- | Reify a runtime 'Int' as a type-level 'GHC.TypeNats.Nat'. +-- +-- The continuation receives a 'Proxy' carrying the reified type. +-- This is useful for working with matrices of runtime-determined size. +reifyDim :: Int -> (forall n. KnownNat n => Proxy n -> a) -> a +reifyDim n f = case someNatVal (fromIntegral n) of + SomeNat p -> f p + +-- | Reify two runtime 'Int's as type-level 'GHC.TypeNats.Nat's. +reifyDim2 :: Int -> Int -> (forall m n. (KnownNat m, KnownNat n) => Proxy m -> Proxy n -> a) -> a +reifyDim2 m n f = reifyDim m $ \pm -> reifyDim n $ \pn -> f pm pn + +-- | Create a matrix using a pure indexing function (sequential computation). +-- +-- @ +-- makeMatrix \@3 \@3 \@P $ \\i j -> fromIntegral (i * 3 + j) +-- @ +-- +-- Internally uses massiv's @'Data.Massiv.Array.Delayed'@ representation as +-- an intermediate before computing to the target representation @r@. +makeMatrix :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e) + => (Int -> Int -> e) -> Matrix m n r e +makeMatrix f = + let r = dimVal @m + c = dimVal @n + in MkMatrix $ M.compute @r $ M.makeArray @D Seq (M.Sz2 r c) (\(i :. j) -> f i j) + +-- | Create a vector using a pure indexing function (sequential computation). +makeVector :: forall n r e. (KnownNat n, M.Manifest r e) + => (Int -> e) -> Vector n r e +makeVector f = + let sz = dimVal @n + in MkVector $ M.compute @r $ M.makeArray @D Seq (M.Sz1 sz) f + +-- | Create a matrix using a pure indexing function with specified +-- 'Data.Massiv.Array.Comp' strategy. +-- +-- Use @Par@ for parallel construction of large matrices: +-- +-- @ +-- makeMatrixComp \@1000 \@1000 \@P Par $ \\i j -> ... +-- @ +makeMatrixComp :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e) + => Comp -> (Int -> Int -> e) -> Matrix m n r e +makeMatrixComp comp f = + let r = dimVal @m + c = dimVal @n + in MkMatrix $ M.compute @r $ M.makeArray @D comp (M.Sz2 r c) (\(i :. j) -> f i j) + +-- | Create a vector using a pure indexing function with specified 'Comp'. +makeVectorComp :: forall n r e. (KnownNat n, M.Manifest r e) + => Comp -> (Int -> e) -> Vector n r e +makeVectorComp comp f = + let sz = dimVal @n + in MkVector $ M.compute @r $ M.makeArray @D comp (M.Sz1 sz) f + +-- | Create a matrix using a mutable 'ST' computation. +-- +-- The action receives a pre-allocated mutable array of the correct size. +-- All writes must be within bounds. The mutable array is frozen after +-- the action completes. +-- +-- This is the primary mechanism for implementing algorithms with in-place +-- updates (e.g., LU factorization, Cholesky decomposition). +createMatrix :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e) + => (forall s. M.MArray s r Ix2 e -> ST s ()) -> Matrix m n r e +createMatrix action = + let r = dimVal @m + c = dimVal @n + arr = M.createArrayST_ (M.Sz2 r c) action + in MkMatrix arr + +-- | Create a vector using a mutable 'ST' computation. +createVector :: forall n r e. (KnownNat n, M.Manifest r e) + => (forall s. M.MArray s r Ix1 e -> ST s ()) -> Vector n r e +createVector action = + let sz = dimVal @n + arr = M.createArrayST_ (M.Sz1 sz) action + in MkVector arr + +-- | Create a matrix using a mutable computation with specified 'Comp'. +createMatrixComp :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e) + => Comp -> (forall s. M.MArray s r Ix2 e -> ST s ()) -> Matrix m n r e +createMatrixComp _comp action = + -- Note: createArrayST_ always runs sequentially; Comp is for delayed computations + createMatrix @m @n action + +-- | Create a vector using a mutable computation with specified 'Comp'. +createVectorComp :: forall n r e. (KnownNat n, M.Manifest r e) + => Comp -> (forall s. M.MArray s r Ix1 e -> ST s ()) -> Vector n r e +createVectorComp _comp action = createVector @n action + +-- | Run a mutable operation on a /copy/ of the matrix, returning both the +-- action's result and the modified matrix. The original matrix is not modified. +withMutableMatrix :: (M.Manifest r e) + => Matrix m n r e + -> (forall s. M.MArray s r Ix2 e -> ST s a) + -> (a, Matrix m n r e) +withMutableMatrix (MkMatrix arr) action = + let (result, arr') = M.withMArrayST arr action + in (result, MkMatrix arr') + +-- | Run a mutable operation on a /copy/ of the vector. +withMutableVector :: (M.Manifest r e) + => Vector n r e + -> (forall s. M.MArray s r Ix1 e -> ST s a) + -> (a, Vector n r e) +withMutableVector (MkVector arr) action = + let (result, arr') = M.withMArrayST arr action + in (result, MkVector arr') + +-- | Like 'withMutableMatrix' but discards the action's result. +withMutableMatrix_ :: (M.Manifest r e) + => Matrix m n r e + -> (forall s. M.MArray s r Ix2 e -> ST s ()) + -> Matrix m n r e +withMutableMatrix_ mat action = snd $ withMutableMatrix mat action + +-- | Like 'withMutableVector' but discards the action's result. +withMutableVector_ :: (M.Manifest r e) + => Vector n r e + -> (forall s. M.MArray s r Ix1 e -> ST s ()) + -> Vector n r e +withMutableVector_ vec action = snd $ withMutableVector vec action + +-- | The \(n \times n\) identity matrix \(I_n\). +-- +-- \[ +-- I_{ij} = \begin{cases} 1 & \text{if } i = j \\ 0 & \text{otherwise} \end{cases} +-- \] +identityMatrix :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Matrix n n r e +identityMatrix = makeMatrix @n @n @r $ \i j -> if i == j then 1 else 0 + +-- | The \(m \times n\) zero matrix. +zeroMatrix :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => Matrix m n r e +zeroMatrix = makeMatrix @m @n @r $ \_ _ -> 0 + +-- | The zero vector of dimension \(n\). +zeroVector :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Vector n r e +zeroVector = makeVector @n @r $ const 0 diff --git a/src/Numeric/LinearAlgebra/Massiv/Internal/Kernel.hs b/src/Numeric/LinearAlgebra/Massiv/Internal/Kernel.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Internal/Kernel.hs @@ -0,0 +1,2440 @@ +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE BangPatterns #-} + +-- | Raw ByteArray# + AVX2 SIMD kernels for hot inner loops. +-- +-- These functions bypass massiv's per-element abstraction layer +-- (M.readM / M.write_ / mapM_) and operate directly on the underlying +-- ByteArray# / MutableByteArray# storage, using GHC's DoubleX4# primops +-- for 256-bit AVX2 SIMD where possible. +module Numeric.LinearAlgebra.Massiv.Internal.Kernel + ( -- * BLAS-1: dot product + rawDot + -- * BLAS-2: matrix-vector multiply + , rawGemv + -- * BLAS-3: matrix multiply (GEMM) + , rawGemmKernel + , rawGemmBISlice + , rawGemmBIBJSlice + , rawGemmBISlicePackedBK + -- * BLAS-3: symmetric rank-k update (SYRK) + , rawSyrkLowerKernel + -- * QR helpers (immutable) + , rawSumSqRange + , rawSumProdRange + , rawHouseholderApplyCol + , rawQAccumCol + -- * Eigen / tridiag helpers + , rawApplyGivensRows + , rawSymRank2Update + -- * LU kernels + , rawLUEliminateColumn + , rawLUEliminateColumnTo + , rawSwapRows + , rawPivotSearch + , rawForwardSubUnitPacked + , rawBackSubPacked + , rawForwardSubUnitPackedSIMD + , rawBackSubPackedSIMD + -- * Cholesky kernels + , rawCholColumn + , rawCholColumnSIMD + , rawCholColumnSIMDFrom + , rawForwardSubCholPacked + , rawBackSubCholTPacked + , rawForwardSubCholPackedSIMD + , rawBackSubCholTPackedSIMD + -- * QR mutable kernels + , rawMutSumSqColumn + , rawMutSumProdColumns + , rawMutHouseholderApply + , rawMutQAccum + -- * Tridiagonalisation mutable kernels + , rawMutSymMatvecSub + , rawMutSymRank2Update + , rawMutTridiagQAccum + -- * Eigen mutable kernels + , rawMutApplyGivensColumns + , rawMutApplyGivensColumnsCM + -- * Matrix transpose + , rawTransposeToColMajor + , rawTransposeFromColMajor + -- * Bulk memory operations + , rawZeroDoubles + , rawCopyDoubles + , rawNegateDoubles + , rawCopyColumn + -- * SVD / bidiagonalisation kernels + , rawMutHouseholderApplyRow + , rawMutSumSqRow + ) where + +import GHC.Exts +import GHC.Prim +import GHC.ST (ST(..)) +import GHC.Types (Double(..), Int(..)) +import Data.Primitive.ByteArray (ByteArray(..), MutableByteArray(..)) + +-- -------------------------------------------------------------------------- +-- BLAS-1: dot product (DoubleX4# FMA, scalar cleanup) +-- -------------------------------------------------------------------------- + +-- | Dot product of two Double vectors stored in ByteArrays. +-- @rawDot ba1 off1 ba2 off2 n@ computes Σ ba1[off1+i] * ba2[off2+i] for i in [0..n-1]. +rawDot :: ByteArray -> Int -> ByteArray -> Int -> Int -> Double +rawDot (ByteArray ba1) (I# off1) (ByteArray ba2) (I# off2) (I# n) = + D# (rawDot# ba1 off1 ba2 off2 n) +{-# INLINE rawDot #-} + +rawDot# :: ByteArray# -> Int# -> ByteArray# -> Int# -> Int# -> Double# +rawDot# ba1 off1 ba2 off2 n = + let n4 = n -# (n `remInt#` 4#) + -- SIMD phase: accumulate 4 doubles at a time + goSimd i acc + | isTrue# (i >=# n4) = acc + | otherwise = + let va = indexDoubleArrayAsDoubleX4# ba1 (off1 +# i) + vb = indexDoubleArrayAsDoubleX4# ba2 (off2 +# i) + in goSimd (i +# 4#) (fmaddDoubleX4# va vb acc) + acc4 = goSimd 0# (broadcastDoubleX4# 0.0##) + !(# a, b, c, d #) = unpackDoubleX4# acc4 + simdSum = a +## b +## c +## d + -- Scalar cleanup for remainder + goScalar i acc + | isTrue# (i >=# n) = acc + | otherwise = + let x = indexDoubleArray# ba1 (off1 +# i) + y = indexDoubleArray# ba2 (off2 +# i) + in goScalar (i +# 1#) (acc +## x *## y) + in goScalar n4 simdSum +{-# NOINLINE rawDot# #-} + +-- -------------------------------------------------------------------------- +-- BLAS-2: matrix-vector multiply +-- -------------------------------------------------------------------------- + +-- | @rawGemv ba_a off_a n_cols ba_x off_x mba_y off_y n_rows@ computes +-- y[i] = Σ_j A[i,j] * x[j] for i in [0..n_rows-1], j in [0..n_cols-1]. +-- A is row-major with stride = n_cols. +rawGemv :: ByteArray -> Int -> Int + -> ByteArray -> Int + -> MutableByteArray s -> Int -> Int + -> ST s () +rawGemv (ByteArray ba_a) (I# off_a) (I# ncols) + (ByteArray ba_x) (I# off_x) + (MutableByteArray mba_y) (I# off_y) (I# nrows) = ST $ \s0 -> + let go i s + | isTrue# (i >=# nrows) = s + | otherwise = + let rowOff = off_a +# i *# ncols + dot = rawDot# ba_a rowOff ba_x off_x ncols + in case writeDoubleArray# mba_y (off_y +# i) dot s of + s' -> go (i +# 1#) s' + in (# go 0# s0, () #) +{-# INLINE rawGemv #-} + +-- -------------------------------------------------------------------------- +-- BLAS-3: tiled ikj GEMM kernel +-- -------------------------------------------------------------------------- + +-- | @rawGemmKernel ba_a off_a ba_b off_b mba_c off_c m k n@ computes +-- C += A * B where A is m×k, B is k×n, C is m×n (all row-major). +-- C must be pre-initialised (e.g. to zero, or to β*C for gemm). +-- Uses packed-B variant (BK-outer with panel packing) for large matrices +-- where cache/TLB effects dominate, unpacked variant for smaller matrices. +rawGemmKernel :: ByteArray -> Int -> ByteArray -> Int + -> MutableByteArray s -> Int + -> Int -> Int -> Int -> ST s () +rawGemmKernel ba offA bb offB mc offC m k n + | min m (min k n) >= gemmPackCrossover = + rawGemmBISlicePackedBK ba offA bb offB mc offC 0 m m k n + | otherwise = + rawGemmBISlice ba offA bb offB mc offC 0 m m k n +{-# INLINE rawGemmKernel #-} + +-- | Crossover threshold for GEMM packing. Below this, the unpacked +-- BI-outer kernel is used; above it, the packed-B BK-outer kernel. +gemmPackCrossover :: Int +gemmPackCrossover = 96 +{-# INLINE gemmPackCrossover #-} + +-- | @rawGemmBISlice@ computes C[biStart..biEnd-1, :] += A[biStart..biEnd-1, :] * B. +-- Delegates to 'rawGemmBIBJSlice' with full column range. +rawGemmBISlice :: ByteArray -> Int -> ByteArray -> Int + -> MutableByteArray s -> Int + -> Int -> Int -> Int -> Int -> Int -> ST s () +rawGemmBISlice ba offA bb offB mc offC biStart biEnd m k n = + rawGemmBIBJSlice ba offA bb offB mc offC biStart biEnd 0 n m k n +{-# INLINE rawGemmBISlice #-} + +-- | @rawGemmBIBJSlice ba_a off_a ba_b off_b mba_c off_c biStart biEnd bjStart bjEnd m k n@ +-- computes C[biStart..biEnd-1, bjStart..bjEnd-1] += A[biStart..biEnd-1, :] * B[:, bjStart..bjEnd-1] +-- where A is m×k, B is k×n, C is m×n (all row-major). +-- Only the rows [biStart, biEnd) and columns [bjStart, bjEnd) of C are written. +-- This enables 2D parallel GEMM by partitioning both row and column ranges. +rawGemmBIBJSlice :: ByteArray -> Int -> ByteArray -> Int + -> MutableByteArray s -> Int + -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> ST s () +rawGemmBIBJSlice (ByteArray ba_a) (I# off_a) (ByteArray ba_b) (I# off_b) + (MutableByteArray mba_c) (I# off_c) + (I# biStart) (I# biEnd) (I# bjStart) (I# bjEnd) + (I# _m) (I# k) (I# n) = ST $ \s0 -> + let bs = 64# + + goBI bi s + | isTrue# (bi >=# biEnd) = s + | otherwise = + let iEnd = minI (bi +# bs) biEnd + in goBI (bi +# bs) (goBK bi iEnd 0# s) + + goBK bi iEnd bk s + | isTrue# (bk >=# k) = s + | otherwise = + let kEnd = minI (bk +# bs) k + in goBK bi iEnd (bk +# bs) (goBJ bi iEnd bk kEnd bjStart s) + + goBJ bi iEnd bk kEnd bj s + | isTrue# (bj >=# bjEnd) = s + | otherwise = + let jEnd = minI (bj +# bs) bjEnd + in goBJ bi iEnd bk kEnd (bj +# bs) (innerBlock bi iEnd bk kEnd bj jEnd s) + + -- Register-blocked micro-kernel: process 4 rows × 8 columns of C + -- in SIMD registers across the full k-range, writing back once. + innerBlock bi iEnd bk kEnd bj jEnd s0_ = + let !jSpan = jEnd -# bj + !j8End = bj +# (jSpan -# (jSpan `remInt#` 8#)) + !j4End = bj +# (jSpan -# (jSpan `remInt#` 4#)) + !iSpan = iEnd -# bi + !i4End = bi +# (iSpan -# (iSpan `remInt#` 4#)) + in goI4 bi i4End iEnd j8End j4End bk kEnd bj jEnd s0_ + + -- Process 4 rows at a time + goI4 i i4End iEnd_ j8End j4End bk kEnd bj jEnd s + | isTrue# (i >=# i4End) = goI1 i iEnd_ j8End j4End bk kEnd bj jEnd s + | otherwise = + goI4 (i +# 4#) i4End iEnd_ j8End j4End bk kEnd bj jEnd + (goJ8_4x8 i bk kEnd bj j8End + (goJ4_4x4 i bk kEnd j8End j4End + (goJScalar4 i bk kEnd j4End jEnd s))) + + -- Process remaining 1 row at a time (up to tile boundary iEnd_, not biEnd) + goI1 i iEnd_ j8End j4End bk kEnd bj jEnd s + | isTrue# (i >=# iEnd_) = s + | otherwise = + goI1 (i +# 1#) iEnd_ j8End j4End bk kEnd bj jEnd + (goJ8_1x8 i bk kEnd bj j8End + (goJ4_1x4 i bk kEnd j8End j4End + (goJScalar1 i bk kEnd j4End jEnd s))) + + -- 4×8 micro-kernel: 4 rows of i, 8 columns of j (2× DoubleX4#) + -- Load 8 C accumulators, sweep k, write back + goJ8_4x8 i bk kEnd j j8End s + | isTrue# (j >=# j8End) = s + | otherwise = + let !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff0 s of { (# s0a, c00 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff0 +# 4#) s0a of { (# s0b, c01 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff1 s0b of { (# s1a, c10 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff1 +# 4#) s1a of { (# s1b, c11 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff2 s1b of { (# s2a, c20 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff2 +# 4#) s2a of { (# s2b, c21 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff3 s2b of { (# s3a, c30 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff3 +# 4#) s3a of { (# s3b, c31 #) -> + -- Now sweep k, accumulating in registers + case goK4x8 i j bk kEnd c00 c01 c10 c11 c20 c21 c30 c31 of + (# r00, r01, r10, r11, r20, r21, r30, r31 #) -> + -- Write back all 8 SIMD registers + case writeDoubleArrayAsDoubleX4# mba_c cOff0 r00 s3b of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff0 +# 4#) r01 sw0 of { sw1 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff1 r10 sw1 of { sw2 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff1 +# 4#) r11 sw2 of { sw3 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff2 r20 sw3 of { sw4 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff2 +# 4#) r21 sw4 of { sw5 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff3 r30 sw5 of { sw6 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff3 +# 4#) r31 sw6 of { sw7 -> + goJ8_4x8 i bk kEnd (j +# 8#) j8End sw7 + }}}}}}}} + }}}}}}}} + + -- Pure k-loop for 4×8: no state threading needed (immutable A, B reads) + goK4x8 i j kk kEnd c00 c01 c10 c11 c20 c21 c30 c31 + | isTrue# (kk >=# kEnd) = + (# c00, c01, c10, c11, c20, c21, c30, c31 #) + | otherwise = + let !bOff = off_b +# kk *# n +# j + !bv0 = indexDoubleArrayAsDoubleX4# ba_b bOff + !bv1 = indexDoubleArrayAsDoubleX4# ba_b (bOff +# 4#) + !a0 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + !a1 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk)) + !a2 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk)) + !a3 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk)) + in goK4x8 i j (kk +# 1#) kEnd + (fmaddDoubleX4# a0 bv0 c00) (fmaddDoubleX4# a0 bv1 c01) + (fmaddDoubleX4# a1 bv0 c10) (fmaddDoubleX4# a1 bv1 c11) + (fmaddDoubleX4# a2 bv0 c20) (fmaddDoubleX4# a2 bv1 c21) + (fmaddDoubleX4# a3 bv0 c30) (fmaddDoubleX4# a3 bv1 c31) + + -- 4×4 cleanup: 4 rows, 4 columns (1× DoubleX4#) + goJ4_4x4 i bk kEnd j j4End s + | isTrue# (j >=# j4End) = s + | otherwise = + let !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff0 s of { (# s0, c0 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff1 s0 of { (# s1, c1 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff2 s1 of { (# s2, c2 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff3 s2 of { (# s3, c3 #) -> + case goK4x4 i j bk kEnd c0 c1 c2 c3 of + (# r0, r1, r2, r3 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff0 r0 s3 of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff1 r1 sw0 of { sw1 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff2 r2 sw1 of { sw2 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff3 r3 sw2 of { sw3 -> + goJ4_4x4 i bk kEnd (j +# 4#) j4End sw3 + }}}} + }}}} + + goK4x4 i j kk kEnd c0 c1 c2 c3 + | isTrue# (kk >=# kEnd) = (# c0, c1, c2, c3 #) + | otherwise = + let !bv = indexDoubleArrayAsDoubleX4# ba_b (off_b +# kk *# n +# j) + !a0 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + !a1 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk)) + !a2 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk)) + !a3 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk)) + in goK4x4 i j (kk +# 1#) kEnd + (fmaddDoubleX4# a0 bv c0) (fmaddDoubleX4# a1 bv c1) + (fmaddDoubleX4# a2 bv c2) (fmaddDoubleX4# a3 bv c3) + + -- Scalar cleanup for 4 rows (columns not a multiple of 4) + goJScalar4 i bk kEnd j jEnd_ s + | isTrue# (j >=# jEnd_) = s + | otherwise = + let goK_s4 kk acc0 acc1 acc2 acc3 + | isTrue# (kk >=# kEnd) = (# acc0, acc1, acc2, acc3 #) + | otherwise = + let !bkj = indexDoubleArray# ba_b (off_b +# kk *# n +# j) + !a0_ = indexDoubleArray# ba_a (off_a +# i *# k +# kk) + !a1_ = indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk) + !a2_ = indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk) + !a3_ = indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk) + in goK_s4 (kk +# 1#) + (acc0 +## a0_ *## bkj) (acc1 +## a1_ *## bkj) + (acc2 +## a2_ *## bkj) (acc3 +## a3_ *## bkj) + !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case (# goK_s4 bk 0.0## 0.0## 0.0## 0.0## #) of + (# (# d0, d1, d2, d3 #) #) -> + case readDoubleArray# mba_c cOff0 s of { (# s0, v0 #) -> + case writeDoubleArray# mba_c cOff0 (v0 +## d0) s0 of { s1 -> + case readDoubleArray# mba_c cOff1 s1 of { (# s2, v1 #) -> + case writeDoubleArray# mba_c cOff1 (v1 +## d1) s2 of { s3 -> + case readDoubleArray# mba_c cOff2 s3 of { (# s4, v2 #) -> + case writeDoubleArray# mba_c cOff2 (v2 +## d2) s4 of { s5 -> + case readDoubleArray# mba_c cOff3 s5 of { (# s6, v3 #) -> + case writeDoubleArray# mba_c cOff3 (v3 +## d3) s6 of { s7 -> + goJScalar4 i bk kEnd (j +# 1#) jEnd_ s7 + }}}}}}}} + + -- 1×8 micro-kernel: 1 row of i, 8 columns of j + goJ8_1x8 i bk kEnd j j8End s + | isTrue# (j >=# j8End) = s + | otherwise = + let !cOff = off_c +# i *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff s of { (# s0, c0 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff +# 4#) s0 of { (# s1, c1 #) -> + case goK1x8 i j bk kEnd c0 c1 of + (# r0, r1 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff r0 s1 of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff +# 4#) r1 sw0 of { sw1 -> + goJ8_1x8 i bk kEnd (j +# 8#) j8End sw1 + }} + }} + + goK1x8 i j kk kEnd c0 c1 + | isTrue# (kk >=# kEnd) = (# c0, c1 #) + | otherwise = + let !bOff = off_b +# kk *# n +# j + !bv0 = indexDoubleArrayAsDoubleX4# ba_b bOff + !bv1 = indexDoubleArrayAsDoubleX4# ba_b (bOff +# 4#) + !av = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + in goK1x8 i j (kk +# 1#) kEnd + (fmaddDoubleX4# av bv0 c0) (fmaddDoubleX4# av bv1 c1) + + -- 1×4 cleanup + goJ4_1x4 i bk kEnd j j4End s + | isTrue# (j >=# j4End) = s + | otherwise = + let !cOff = off_c +# i *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff s of { (# s0, c0 #) -> + case goK1x4 i j bk kEnd c0 of { r0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff r0 s0 of { s1 -> + goJ4_1x4 i bk kEnd (j +# 4#) j4End s1 + }}} + + goK1x4 i j kk kEnd c0 + | isTrue# (kk >=# kEnd) = c0 + | otherwise = + let !bv = indexDoubleArrayAsDoubleX4# ba_b (off_b +# kk *# n +# j) + !av = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + in goK1x4 i j (kk +# 1#) kEnd (fmaddDoubleX4# av bv c0) + + -- Scalar cleanup for 1 row + goJScalar1 i bk kEnd j jEnd_ s + | isTrue# (j >=# jEnd_) = s + | otherwise = + let goK_s1 kk acc + | isTrue# (kk >=# kEnd) = acc + | otherwise = + let !bkj = indexDoubleArray# ba_b (off_b +# kk *# n +# j) + !aik = indexDoubleArray# ba_a (off_a +# i *# k +# kk) + in goK_s1 (kk +# 1#) (acc +## aik *## bkj) + in case readDoubleArray# mba_c (off_c +# i *# n +# j) s of + (# s', cij #) -> + case writeDoubleArray# mba_c (off_c +# i *# n +# j) (cij +## goK_s1 bk 0.0##) s' of + s'' -> goJScalar1 i bk kEnd (j +# 1#) jEnd_ s'' + + in (# goBI biStart s0, () #) +{-# INLINE rawGemmBIBJSlice #-} + +-- | Packed-B GEMM with BK-outer loop ordering. +-- Packs B into 8-column panels for sequential cache access (stride 8 instead +-- of stride n). Only the rows [biStart, biEnd) of C are written. +-- Beneficial for large matrices where stride-n B access causes cache/TLB misses. +rawGemmBISlicePackedBK :: ByteArray -> Int -> ByteArray -> Int + -> MutableByteArray s -> Int + -> Int -> Int -> Int -> Int -> Int -> ST s () +rawGemmBISlicePackedBK (ByteArray ba_a) (I# off_a) (ByteArray ba_b) (I# off_b) + (MutableByteArray mba_c) (I# off_c) + (I# biStart) (I# biEnd) (I# _m) (I# k) (I# n) = ST $ \s0 -> + let bs = 64# + !nPanels = n `quotInt#` 8# + !j8End = nPanels *# 8# + !j4End = n -# (n `remInt#` 4#) + -- Packed buffer: nPanels * bs * 8 doubles (each panel: kc × 8) + !packDoubles = nPanels *# bs *# 8# + + -- Fallback unpacked path (when j8End == 0, i.e. n < 8) + goBI_unpacked bi s + | isTrue# (bi >=# biEnd) = s + | otherwise = + let iEnd = minI (bi +# bs) biEnd + in goBI_unpacked (bi +# bs) (goBK_u bi iEnd 0# s) + + goBK_u bi iEnd bk s + | isTrue# (bk >=# k) = s + | otherwise = + let kEnd = minI (bk +# bs) k + iSpan = iEnd -# bi + i4End_ = bi +# (iSpan -# (iSpan `remInt#` 4#)) + in goBK_u bi iEnd (bk +# bs) + (goI4U_fb bi i4End_ iEnd bk kEnd + (goI1U_fb i4End_ iEnd bk kEnd s)) + + goI4U_fb i i4End _iEnd bk kEnd s + | isTrue# (i >=# i4End) = s + | otherwise = + goI4U_fb (i +# 4#) i4End _iEnd bk kEnd + (goJ4Ufb4 i bk kEnd 0# j4End + (goJSUfb4 i bk kEnd j4End n s)) + + goI1U_fb i iEnd bk kEnd s + | isTrue# (i >=# iEnd) = s + | otherwise = + goI1U_fb (i +# 1#) iEnd bk kEnd + (goJ4Ufb1 i bk kEnd 0# j4End + (goJSUfb1 i bk kEnd j4End n s)) + + -- Unpacked 4×4 and scalar for fallback + goJ4Ufb4 i bk kEnd j j4EndV s + | isTrue# (j >=# j4EndV) = s + | otherwise = + let !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff0 s of { (# s0, c0 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff1 s0 of { (# s1, c1 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff2 s1 of { (# s2, c2 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff3 s2 of { (# s3, c3 #) -> + case goKUfb4x4 i j bk kEnd c0 c1 c2 c3 of + (# r0, r1, r2, r3 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff0 r0 s3 of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff1 r1 sw0 of { sw1 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff2 r2 sw1 of { sw2 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff3 r3 sw2 of { sw3 -> + goJ4Ufb4 i bk kEnd (j +# 4#) j4EndV sw3 + }}}} + }}}} + + goKUfb4x4 i j kk kEnd c0 c1 c2 c3 + | isTrue# (kk >=# kEnd) = (# c0, c1, c2, c3 #) + | otherwise = + let !bv = indexDoubleArrayAsDoubleX4# ba_b (off_b +# kk *# n +# j) + !a0 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + !a1 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk)) + !a2 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk)) + !a3 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk)) + in goKUfb4x4 i j (kk +# 1#) kEnd + (fmaddDoubleX4# a0 bv c0) (fmaddDoubleX4# a1 bv c1) + (fmaddDoubleX4# a2 bv c2) (fmaddDoubleX4# a3 bv c3) + + goJSUfb4 i bk kEnd j jEndV s + | isTrue# (j >=# jEndV) = s + | otherwise = + let goKSfb4 kk acc0 acc1 acc2 acc3 + | isTrue# (kk >=# kEnd) = (# acc0, acc1, acc2, acc3 #) + | otherwise = + let !bkj = indexDoubleArray# ba_b (off_b +# kk *# n +# j) + !a0_ = indexDoubleArray# ba_a (off_a +# i *# k +# kk) + !a1_ = indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk) + !a2_ = indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk) + !a3_ = indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk) + in goKSfb4 (kk +# 1#) + (acc0 +## a0_ *## bkj) (acc1 +## a1_ *## bkj) + (acc2 +## a2_ *## bkj) (acc3 +## a3_ *## bkj) + !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case (# goKSfb4 bk 0.0## 0.0## 0.0## 0.0## #) of + (# (# d0, d1, d2, d3 #) #) -> + case readDoubleArray# mba_c cOff0 s of { (# s0, v0 #) -> + case writeDoubleArray# mba_c cOff0 (v0 +## d0) s0 of { s1 -> + case readDoubleArray# mba_c cOff1 s1 of { (# s2, v1 #) -> + case writeDoubleArray# mba_c cOff1 (v1 +## d1) s2 of { s3 -> + case readDoubleArray# mba_c cOff2 s3 of { (# s4, v2 #) -> + case writeDoubleArray# mba_c cOff2 (v2 +## d2) s4 of { s5 -> + case readDoubleArray# mba_c cOff3 s5 of { (# s6, v3 #) -> + case writeDoubleArray# mba_c cOff3 (v3 +## d3) s6 of { s7 -> + goJSUfb4 i bk kEnd (j +# 1#) jEndV s7 + }}}}}}}} + + goJ4Ufb1 i bk kEnd j j4EndV s + | isTrue# (j >=# j4EndV) = s + | otherwise = + let !cOff = off_c +# i *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff s of { (# s0, c0 #) -> + case goKUfb1x4 i j bk kEnd c0 of { r0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff r0 s0 of { s1 -> + goJ4Ufb1 i bk kEnd (j +# 4#) j4EndV s1 + }}} + + goKUfb1x4 i j kk kEnd c0 + | isTrue# (kk >=# kEnd) = c0 + | otherwise = + let !bv = indexDoubleArrayAsDoubleX4# ba_b (off_b +# kk *# n +# j) + !av = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + in goKUfb1x4 i j (kk +# 1#) kEnd (fmaddDoubleX4# av bv c0) + + goJSUfb1 i bk kEnd j jEndV s + | isTrue# (j >=# jEndV) = s + | otherwise = + let goKSfb1 kk acc + | isTrue# (kk >=# kEnd) = acc + | otherwise = + let !bkj = indexDoubleArray# ba_b (off_b +# kk *# n +# j) + !aik = indexDoubleArray# ba_a (off_a +# i *# k +# kk) + in goKSfb1 (kk +# 1#) (acc +## aik *## bkj) + in case readDoubleArray# mba_c (off_c +# i *# n +# j) s of + (# s', cij #) -> + case writeDoubleArray# mba_c (off_c +# i *# n +# j) (cij +## goKSfb1 bk 0.0##) s' of + s'' -> goJSUfb1 i bk kEnd (j +# 1#) jEndV s'' + + in if isTrue# (j8End ==# 0#) + -- No full 8-column panels: fall back to unpacked scalar path + then (# goBI_unpacked biStart s0, () #) + else + + -- BK outer loop: allocate pack buffer, pack B, freeze, then run BI/BJ loops + -- with pure k-loop reads (indexDoubleArrayAsDoubleX4# on immutable ByteArray#). + let goBKO bk s + | isTrue# (bk >=# k) = s + | otherwise = + let !kEnd = minI (bk +# bs) k + !kc = kEnd -# bk + !packBufBytes = nPanels *# kc *# 8# *# 8# + in case newByteArray# packBufBytes s of { (# s1, mba_bp #) -> + -- Pack B[bk:bk+kc, 0:j8End] into panel-major layout + let s2 = packB mba_bp bk kc s1 + in case unsafeFreezeByteArray# mba_bp s2 of { (# s3, ba_bp #) -> + goBKO (bk +# bs) (goBIO bk kEnd kc ba_bp biStart s3) + }} + + -- Pack B into mba_bp: panel p (cols 8p..8p+7): + -- bp[p * kc * 8 + kLocal * 8 + jLocal] = B[bk+kLocal, 8p+jLocal] + packB mba_bp bk kc sp = + let goPnl p sp0 + | isTrue# (p >=# nPanels) = sp0 + | otherwise = + let !jBase = p *# 8# + !pOff = p *# kc *# 8# + goKP kl sk + | isTrue# (kl >=# kc) = sk + | otherwise = + let !srcOff = off_b +# (bk +# kl) *# n +# jBase + !dstOff = pOff +# kl *# 8# + !v0 = indexDoubleArrayAsDoubleX4# ba_b srcOff + !v1 = indexDoubleArrayAsDoubleX4# ba_b (srcOff +# 4#) + in case writeDoubleArrayAsDoubleX4# mba_bp dstOff v0 sk of + sk' -> case writeDoubleArrayAsDoubleX4# mba_bp (dstOff +# 4#) v1 sk' of + sk'' -> goKP (kl +# 1#) sk'' + in goPnl (p +# 1#) (goKP 0# sp0) + in goPnl 0# sp + + -- BI middle loop (ba_bp is frozen ByteArray# for this BK tile) + goBIO bk kEnd kc ba_bp bi s + | isTrue# (bi >=# biEnd) = s + | otherwise = + let !iEnd = minI (bi +# bs) biEnd + !iSpan = iEnd -# bi + !i4End_ = bi +# (iSpan -# (iSpan `remInt#` 4#)) + in goBIO bk kEnd kc ba_bp (bi +# bs) + (goI4P bi i4End_ iEnd bk kEnd kc ba_bp + (goI1P i4End_ iEnd bk kEnd kc ba_bp s)) + + -- 4 rows: packed panels then unpacked tail + goI4P i i4End _iEnd bk kEnd kc ba_bp s + | isTrue# (i >=# i4End) = s + | otherwise = + goI4P (i +# 4#) i4End _iEnd bk kEnd kc ba_bp + (goJ8P4 i bk kc ba_bp 0# j8End + (goJ4U4 i bk kEnd j8End j4End + (goJSU4 i bk kEnd j4End n s))) + + -- 1 row: packed panels then unpacked tail + goI1P i iEnd bk kEnd kc ba_bp s + | isTrue# (i >=# iEnd) = s + | otherwise = + goI1P (i +# 1#) iEnd bk kEnd kc ba_bp + (goJ8P1 i bk kc ba_bp 0# j8End + (goJ4U1 i bk kEnd j8End j4End + (goJSU1 i bk kEnd j4End n s))) + + -- ---------------------------------------------------------------- + -- Packed 4×8 j-loop: step by 8, pure reads from frozen ByteArray# + -- ---------------------------------------------------------------- + goJ8P4 i bk kc ba_bp j jEnd s + | isTrue# (j >=# jEnd) = s + | otherwise = + let !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + !panOff = (j `quotInt#` 8#) *# kc *# 8# + in case readDoubleArrayAsDoubleX4# mba_c cOff0 s of { (# s0a, c00 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff0 +# 4#) s0a of { (# s0b, c01 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff1 s0b of { (# s1a, c10 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff1 +# 4#) s1a of { (# s1b, c11 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff2 s1b of { (# s2a, c20 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff2 +# 4#) s2a of { (# s2b, c21 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff3 s2b of { (# s3a, c30 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff3 +# 4#) s3a of { (# s3b, c31 #) -> + case goKP4x8 i panOff bk 0# kc ba_bp c00 c01 c10 c11 c20 c21 c30 c31 of + (# r00, r01, r10, r11, r20, r21, r30, r31 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff0 r00 s3b of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff0 +# 4#) r01 sw0 of { sw1 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff1 r10 sw1 of { sw2 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff1 +# 4#) r11 sw2 of { sw3 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff2 r20 sw3 of { sw4 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff2 +# 4#) r21 sw4 of { sw5 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff3 r30 sw5 of { sw6 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff3 +# 4#) r31 sw6 of { sw7 -> + goJ8P4 i bk kc ba_bp (j +# 8#) jEnd sw7 + }}}}}}}} + }}}}}}}} + + -- Pure k-loop for packed 4×8: reads from frozen ByteArray# (no State#) + goKP4x8 i panOff bk kl kc ba_bp c00 c01 c10 c11 c20 c21 c30 c31 + | isTrue# (kl >=# kc) = + (# c00, c01, c10, c11, c20, c21, c30, c31 #) + | otherwise = + let !bOff = panOff +# kl *# 8# + !bv0 = indexDoubleArrayAsDoubleX4# ba_bp bOff + !bv1 = indexDoubleArrayAsDoubleX4# ba_bp (bOff +# 4#) + !kk = bk +# kl + !a0 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + !a1 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk)) + !a2 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk)) + !a3 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk)) + in goKP4x8 i panOff bk (kl +# 1#) kc ba_bp + (fmaddDoubleX4# a0 bv0 c00) (fmaddDoubleX4# a0 bv1 c01) + (fmaddDoubleX4# a1 bv0 c10) (fmaddDoubleX4# a1 bv1 c11) + (fmaddDoubleX4# a2 bv0 c20) (fmaddDoubleX4# a2 bv1 c21) + (fmaddDoubleX4# a3 bv0 c30) (fmaddDoubleX4# a3 bv1 c31) + + -- ---------------------------------------------------------------- + -- Packed 1×8 j-loop + -- ---------------------------------------------------------------- + goJ8P1 i bk kc ba_bp j jEnd s + | isTrue# (j >=# jEnd) = s + | otherwise = + let !cOff = off_c +# i *# n +# j + !panOff = (j `quotInt#` 8#) *# kc *# 8# + in case readDoubleArrayAsDoubleX4# mba_c cOff s of { (# s0, c0 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff +# 4#) s0 of { (# s1, c1 #) -> + case goKP1x8 i panOff bk 0# kc ba_bp c0 c1 of + (# r0, r1 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff r0 s1 of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff +# 4#) r1 sw0 of { sw1 -> + goJ8P1 i bk kc ba_bp (j +# 8#) jEnd sw1 + }} + }} + + -- Pure k-loop for packed 1×8 + goKP1x8 i panOff bk kl kc ba_bp c0 c1 + | isTrue# (kl >=# kc) = (# c0, c1 #) + | otherwise = + let !bOff = panOff +# kl *# 8# + !bv0 = indexDoubleArrayAsDoubleX4# ba_bp bOff + !bv1 = indexDoubleArrayAsDoubleX4# ba_bp (bOff +# 4#) + !kk = bk +# kl + !av = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + in goKP1x8 i panOff bk (kl +# 1#) kc ba_bp + (fmaddDoubleX4# av bv0 c0) (fmaddDoubleX4# av bv1 c1) + + -- ---------------------------------------------------------------- + -- Unpacked j-tail: 4×4 cleanup (columns j8End..j4End) + -- ---------------------------------------------------------------- + goJ4U4 i bk kEnd j j4EndV s + | isTrue# (j >=# j4EndV) = s + | otherwise = + let !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff0 s of { (# s0, c0 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff1 s0 of { (# s1, c1 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff2 s1 of { (# s2, c2 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff3 s2 of { (# s3, c3 #) -> + case goKU4x4 i j bk kEnd c0 c1 c2 c3 of + (# r0, r1, r2, r3 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff0 r0 s3 of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff1 r1 sw0 of { sw1 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff2 r2 sw1 of { sw2 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff3 r3 sw2 of { sw3 -> + goJ4U4 i bk kEnd (j +# 4#) j4EndV sw3 + }}}} + }}}} + + goKU4x4 i j kk kEnd c0 c1 c2 c3 + | isTrue# (kk >=# kEnd) = (# c0, c1, c2, c3 #) + | otherwise = + let !bv = indexDoubleArrayAsDoubleX4# ba_b (off_b +# kk *# n +# j) + !a0 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + !a1 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk)) + !a2 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk)) + !a3 = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk)) + in goKU4x4 i j (kk +# 1#) kEnd + (fmaddDoubleX4# a0 bv c0) (fmaddDoubleX4# a1 bv c1) + (fmaddDoubleX4# a2 bv c2) (fmaddDoubleX4# a3 bv c3) + + -- Unpacked scalar cleanup for 4 rows (columns j4End..n) + goJSU4 i bk kEnd j jEndV s + | isTrue# (j >=# jEndV) = s + | otherwise = + let goKS4 kk acc0 acc1 acc2 acc3 + | isTrue# (kk >=# kEnd) = (# acc0, acc1, acc2, acc3 #) + | otherwise = + let !bkj = indexDoubleArray# ba_b (off_b +# kk *# n +# j) + !a0_ = indexDoubleArray# ba_a (off_a +# i *# k +# kk) + !a1_ = indexDoubleArray# ba_a (off_a +# (i +# 1#) *# k +# kk) + !a2_ = indexDoubleArray# ba_a (off_a +# (i +# 2#) *# k +# kk) + !a3_ = indexDoubleArray# ba_a (off_a +# (i +# 3#) *# k +# kk) + in goKS4 (kk +# 1#) + (acc0 +## a0_ *## bkj) (acc1 +## a1_ *## bkj) + (acc2 +## a2_ *## bkj) (acc3 +## a3_ *## bkj) + !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case (# goKS4 bk 0.0## 0.0## 0.0## 0.0## #) of + (# (# d0, d1, d2, d3 #) #) -> + case readDoubleArray# mba_c cOff0 s of { (# s0, v0 #) -> + case writeDoubleArray# mba_c cOff0 (v0 +## d0) s0 of { s1 -> + case readDoubleArray# mba_c cOff1 s1 of { (# s2, v1 #) -> + case writeDoubleArray# mba_c cOff1 (v1 +## d1) s2 of { s3 -> + case readDoubleArray# mba_c cOff2 s3 of { (# s4, v2 #) -> + case writeDoubleArray# mba_c cOff2 (v2 +## d2) s4 of { s5 -> + case readDoubleArray# mba_c cOff3 s5 of { (# s6, v3 #) -> + case writeDoubleArray# mba_c cOff3 (v3 +## d3) s6 of { s7 -> + goJSU4 i bk kEnd (j +# 1#) jEndV s7 + }}}}}}}} + + -- Unpacked 1×4 cleanup + goJ4U1 i bk kEnd j j4EndV s + | isTrue# (j >=# j4EndV) = s + | otherwise = + let !cOff = off_c +# i *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff s of { (# s0, c0 #) -> + case goKU1x4 i j bk kEnd c0 of { r0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff r0 s0 of { s1 -> + goJ4U1 i bk kEnd (j +# 4#) j4EndV s1 + }}} + + goKU1x4 i j kk kEnd c0 + | isTrue# (kk >=# kEnd) = c0 + | otherwise = + let !bv = indexDoubleArrayAsDoubleX4# ba_b (off_b +# kk *# n +# j) + !av = broadcastDoubleX4# (indexDoubleArray# ba_a (off_a +# i *# k +# kk)) + in goKU1x4 i j (kk +# 1#) kEnd (fmaddDoubleX4# av bv c0) + + -- Unpacked scalar cleanup for 1 row + goJSU1 i bk kEnd j jEndV s + | isTrue# (j >=# jEndV) = s + | otherwise = + let goKS1 kk acc + | isTrue# (kk >=# kEnd) = acc + | otherwise = + let !bkj = indexDoubleArray# ba_b (off_b +# kk *# n +# j) + !aik = indexDoubleArray# ba_a (off_a +# i *# k +# kk) + in goKS1 (kk +# 1#) (acc +## aik *## bkj) + in case readDoubleArray# mba_c (off_c +# i *# n +# j) s of + (# s', cij #) -> + case writeDoubleArray# mba_c (off_c +# i *# n +# j) (cij +## goKS1 bk 0.0##) s' of + s'' -> goJSU1 i bk kEnd (j +# 1#) jEndV s'' + + in (# goBKO 0# s0, () #) +{-# INLINE rawGemmBISlicePackedBK #-} + +-- | @rawSyrkLowerKernel ba_a off_a mba_c off_c m n@ +-- computes C = A^T * A where A is m×n (row-major). +-- Only the lower triangle of C (n×n) is computed via tiled SIMD micro-kernels, +-- then mirrored to the upper triangle. C must be pre-zeroed. +-- This avoids materialising A^T and halves the flop count vs full GEMM. +rawSyrkLowerKernel :: ByteArray -> Int + -> MutableByteArray s -> Int + -> Int -> Int -> ST s () +rawSyrkLowerKernel (ByteArray ba_a) (I# off_a) + (MutableByteArray mba_c) (I# off_c) + (I# m) (I# n) = ST $ \s0 -> + let bs = 64# + + -- Tile rows of C (i dimension) + goBI bi s + | isTrue# (bi >=# n) = s + | otherwise = + let iEnd = minI (bi +# bs) n + in goBI (bi +# bs) (goBK bi iEnd 0# s) + + -- Tile inner dimension (k = rows of A, 0..m-1) + goBK bi iEnd bk s + | isTrue# (bk >=# m) = s + | otherwise = + let kEnd = minI (bk +# bs) m + in goBK bi iEnd (bk +# bs) (goBJ bi iEnd bk kEnd 0# s) + + -- Tile columns of C (j dimension) — lower triangle only: bj <= bi + goBJ bi iEnd bk kEnd bj s + | isTrue# (bj ># bi) = s -- stop past diagonal + | otherwise = + let jEnd = minI (bj +# bs) n + in goBJ bi iEnd bk kEnd (bj +# bs) (innerBlock bi iEnd bk kEnd bj jEnd s) + + -- Micro-kernel dispatch (same structure as GEMM) + innerBlock bi iEnd bk kEnd bj jEnd s0_ = + let !jSpan = jEnd -# bj + !j8End = bj +# (jSpan -# (jSpan `remInt#` 8#)) + !j4End = bj +# (jSpan -# (jSpan `remInt#` 4#)) + !iSpan = iEnd -# bi + !i4End = bi +# (iSpan -# (iSpan `remInt#` 4#)) + in goI4 bi i4End iEnd j8End j4End bk kEnd bj jEnd s0_ + + goI4 i i4End iEnd_ j8End j4End bk kEnd bj jEnd s + | isTrue# (i >=# i4End) = goI1 i iEnd_ j8End j4End bk kEnd bj jEnd s + | otherwise = + goI4 (i +# 4#) i4End iEnd_ j8End j4End bk kEnd bj jEnd + (goJ8s i bk kEnd bj j8End + (goJ4s i bk kEnd j8End j4End + (goJSs4 i bk kEnd j4End jEnd s))) + + goI1 i iEnd_ j8End j4End bk kEnd bj jEnd s + | isTrue# (i >=# iEnd_) = s + | otherwise = + goI1 (i +# 1#) iEnd_ j8End j4End bk kEnd bj jEnd + (goJ8s1 i bk kEnd bj j8End + (goJ4s1 i bk kEnd j8End j4End + (goJSs1 i bk kEnd j4End jEnd s))) + + -- 4×8 micro-kernel for SYRK + -- A^T[i,kk] = A[kk,i] at off_a + kk*n + i + -- A[kk,j] at off_a + kk*n + j + goJ8s i bk kEnd j j8End s + | isTrue# (j >=# j8End) = s + | otherwise = + let !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff0 s of { (# s0a, c00 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff0 +# 4#) s0a of { (# s0b, c01 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff1 s0b of { (# s1a, c10 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff1 +# 4#) s1a of { (# s1b, c11 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff2 s1b of { (# s2a, c20 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff2 +# 4#) s2a of { (# s2b, c21 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff3 s2b of { (# s3a, c30 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff3 +# 4#) s3a of { (# s3b, c31 #) -> + case goK8s i j bk kEnd c00 c01 c10 c11 c20 c21 c30 c31 of + (# r00, r01, r10, r11, r20, r21, r30, r31 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff0 r00 s3b of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff0 +# 4#) r01 sw0 of { sw1 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff1 r10 sw1 of { sw2 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff1 +# 4#) r11 sw2 of { sw3 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff2 r20 sw3 of { sw4 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff2 +# 4#) r21 sw4 of { sw5 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff3 r30 sw5 of { sw6 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff3 +# 4#) r31 sw6 of { sw7 -> + goJ8s i bk kEnd (j +# 8#) j8End sw7 + }}}}}}}} + }}}}}}}} + + -- k-loop for 4×8 SYRK: A^T[i,kk]=A[kk,i], A[kk,j..j+7] + goK8s i j kk kEnd c00 c01 c10 c11 c20 c21 c30 c31 + | isTrue# (kk >=# kEnd) = + (# c00, c01, c10, c11, c20, c21, c30, c31 #) + | otherwise = + let !rowOff = off_a +# kk *# n + !bv0 = indexDoubleArrayAsDoubleX4# ba_a (rowOff +# j) + !bv1 = indexDoubleArrayAsDoubleX4# ba_a (rowOff +# j +# 4#) + !a0 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i)) + !a1 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i +# 1#)) + !a2 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i +# 2#)) + !a3 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i +# 3#)) + in goK8s i j (kk +# 1#) kEnd + (fmaddDoubleX4# a0 bv0 c00) (fmaddDoubleX4# a0 bv1 c01) + (fmaddDoubleX4# a1 bv0 c10) (fmaddDoubleX4# a1 bv1 c11) + (fmaddDoubleX4# a2 bv0 c20) (fmaddDoubleX4# a2 bv1 c21) + (fmaddDoubleX4# a3 bv0 c30) (fmaddDoubleX4# a3 bv1 c31) + + -- 4×4 cleanup for SYRK + goJ4s i bk kEnd j j4End s + | isTrue# (j >=# j4End) = s + | otherwise = + let !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff0 s of { (# s0, c0 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff1 s0 of { (# s1, c1 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff2 s1 of { (# s2, c2 #) -> + case readDoubleArrayAsDoubleX4# mba_c cOff3 s2 of { (# s3, c3 #) -> + case goK4s i j bk kEnd c0 c1 c2 c3 of + (# r0, r1, r2, r3 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff0 r0 s3 of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff1 r1 sw0 of { sw1 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff2 r2 sw1 of { sw2 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff3 r3 sw2 of { sw3 -> + goJ4s i bk kEnd (j +# 4#) j4End sw3 + }}}} + }}}} + + goK4s i j kk kEnd c0 c1 c2 c3 + | isTrue# (kk >=# kEnd) = (# c0, c1, c2, c3 #) + | otherwise = + let !rowOff = off_a +# kk *# n + !bv = indexDoubleArrayAsDoubleX4# ba_a (rowOff +# j) + !a0 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i)) + !a1 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i +# 1#)) + !a2 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i +# 2#)) + !a3 = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i +# 3#)) + in goK4s i j (kk +# 1#) kEnd + (fmaddDoubleX4# a0 bv c0) (fmaddDoubleX4# a1 bv c1) + (fmaddDoubleX4# a2 bv c2) (fmaddDoubleX4# a3 bv c3) + + -- Scalar cleanup for 4 rows (SYRK) + goJSs4 i bk kEnd j jEnd_ s + | isTrue# (j >=# jEnd_) = s + | otherwise = + let goK_s4 kk acc0 acc1 acc2 acc3 + | isTrue# (kk >=# kEnd) = (# acc0, acc1, acc2, acc3 #) + | otherwise = + let !rowOff = off_a +# kk *# n + !bkj = indexDoubleArray# ba_a (rowOff +# j) + !a0_ = indexDoubleArray# ba_a (rowOff +# i) + !a1_ = indexDoubleArray# ba_a (rowOff +# i +# 1#) + !a2_ = indexDoubleArray# ba_a (rowOff +# i +# 2#) + !a3_ = indexDoubleArray# ba_a (rowOff +# i +# 3#) + in goK_s4 (kk +# 1#) + (acc0 +## a0_ *## bkj) (acc1 +## a1_ *## bkj) + (acc2 +## a2_ *## bkj) (acc3 +## a3_ *## bkj) + !cOff0 = off_c +# i *# n +# j + !cOff1 = off_c +# (i +# 1#) *# n +# j + !cOff2 = off_c +# (i +# 2#) *# n +# j + !cOff3 = off_c +# (i +# 3#) *# n +# j + in case (# goK_s4 bk 0.0## 0.0## 0.0## 0.0## #) of + (# (# d0, d1, d2, d3 #) #) -> + case readDoubleArray# mba_c cOff0 s of { (# s0, v0 #) -> + case writeDoubleArray# mba_c cOff0 (v0 +## d0) s0 of { s1 -> + case readDoubleArray# mba_c cOff1 s1 of { (# s2, v1 #) -> + case writeDoubleArray# mba_c cOff1 (v1 +## d1) s2 of { s3 -> + case readDoubleArray# mba_c cOff2 s3 of { (# s4, v2 #) -> + case writeDoubleArray# mba_c cOff2 (v2 +## d2) s4 of { s5 -> + case readDoubleArray# mba_c cOff3 s5 of { (# s6, v3 #) -> + case writeDoubleArray# mba_c cOff3 (v3 +## d3) s6 of { s7 -> + goJSs4 i bk kEnd (j +# 1#) jEnd_ s7 + }}}}}}}} + + -- 1×8 micro-kernel for SYRK + goJ8s1 i bk kEnd j j8End s + | isTrue# (j >=# j8End) = s + | otherwise = + let !cOff = off_c +# i *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff s of { (# s0, c0 #) -> + case readDoubleArrayAsDoubleX4# mba_c (cOff +# 4#) s0 of { (# s1, c1 #) -> + case goK8s1 i j bk kEnd c0 c1 of + (# r0, r1 #) -> + case writeDoubleArrayAsDoubleX4# mba_c cOff r0 s1 of { sw0 -> + case writeDoubleArrayAsDoubleX4# mba_c (cOff +# 4#) r1 sw0 of { sw1 -> + goJ8s1 i bk kEnd (j +# 8#) j8End sw1 + }} + }} + + goK8s1 i j kk kEnd c0 c1 + | isTrue# (kk >=# kEnd) = (# c0, c1 #) + | otherwise = + let !rowOff = off_a +# kk *# n + !bv0 = indexDoubleArrayAsDoubleX4# ba_a (rowOff +# j) + !bv1 = indexDoubleArrayAsDoubleX4# ba_a (rowOff +# j +# 4#) + !av = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i)) + in goK8s1 i j (kk +# 1#) kEnd + (fmaddDoubleX4# av bv0 c0) (fmaddDoubleX4# av bv1 c1) + + -- 1×4 cleanup for SYRK + goJ4s1 i bk kEnd j j4End s + | isTrue# (j >=# j4End) = s + | otherwise = + let !cOff = off_c +# i *# n +# j + in case readDoubleArrayAsDoubleX4# mba_c cOff s of { (# s0, c0 #) -> + case goK4s1 i j bk kEnd c0 of { r0 -> + case writeDoubleArrayAsDoubleX4# mba_c cOff r0 s0 of { s1 -> + goJ4s1 i bk kEnd (j +# 4#) j4End s1 + }}} + + goK4s1 i j kk kEnd c0 + | isTrue# (kk >=# kEnd) = c0 + | otherwise = + let !rowOff = off_a +# kk *# n + !bv = indexDoubleArrayAsDoubleX4# ba_a (rowOff +# j) + !av = broadcastDoubleX4# (indexDoubleArray# ba_a (rowOff +# i)) + in goK4s1 i j (kk +# 1#) kEnd (fmaddDoubleX4# av bv c0) + + -- Scalar cleanup for 1 row (SYRK) + goJSs1 i bk kEnd j jEnd_ s + | isTrue# (j >=# jEnd_) = s + | otherwise = + let goK_s1 kk acc + | isTrue# (kk >=# kEnd) = acc + | otherwise = + let !rowOff = off_a +# kk *# n + !bkj = indexDoubleArray# ba_a (rowOff +# j) + !aik = indexDoubleArray# ba_a (rowOff +# i) + in goK_s1 (kk +# 1#) (acc +## aik *## bkj) + in case readDoubleArray# mba_c (off_c +# i *# n +# j) s of + (# s', cij #) -> + case writeDoubleArray# mba_c (off_c +# i *# n +# j) (cij +## goK_s1 bk 0.0##) s' of + s'' -> goJSs1 i bk kEnd (j +# 1#) jEnd_ s'' + + -- Mirror lower triangle to upper: C[j,i] = C[i,j] for j < i + mirror i s + | isTrue# (i >=# n) = s + | otherwise = mirror (i +# 1#) (mirrorRow i 0# s) + + mirrorRow i j s + | isTrue# (j >=# i) = s + | otherwise = + case readDoubleArray# mba_c (off_c +# i *# n +# j) s of { (# s0, cij #) -> + case writeDoubleArray# mba_c (off_c +# j *# n +# i) cij s0 of { s1 -> + mirrorRow i (j +# 1#) s1 + }} + + in (# mirror 0# (goBI 0# s0), () #) +{-# INLINE rawSyrkLowerKernel #-} + +-- -------------------------------------------------------------------------- +-- QR helpers +-- -------------------------------------------------------------------------- + +-- | Sum of squares: Σ arr[off+i]^2 for i in [from..to-1] +rawSumSqRange :: ByteArray -> Int -> Int -> Int -> Double +rawSumSqRange (ByteArray ba) (I# off) (I# from_) (I# to) = + D# (goSumSq from_ 0.0##) + where + goSumSq i acc + | isTrue# (i >=# to) = acc + | otherwise = + let x = indexDoubleArray# ba (off +# i) + in goSumSq (i +# 1#) (acc +## x *## x) +{-# INLINE rawSumSqRange #-} + +-- | Dot product of a column slice: Σ arr1[off1+i*stride1] * arr2[off2+i*stride2] +-- for i in [from..to-1]. Used for column-wise access patterns in QR. +rawSumProdRange :: ByteArray -> Int -> Int + -> ByteArray -> Int -> Int + -> Int -> Int -> Double +rawSumProdRange (ByteArray ba1) (I# off1) (I# stride1) + (ByteArray ba2) (I# off2) (I# stride2) + (I# from_) (I# to) = + D# (goSumProd from_ 0.0##) + where + goSumProd i acc + | isTrue# (i >=# to) = acc + | otherwise = + let x = indexDoubleArray# ba1 (off1 +# i *# stride1) + y = indexDoubleArray# ba2 (off2 +# i *# stride2) + in goSumProd (i +# 1#) (acc +## x *## y) +{-# INLINE rawSumProdRange #-} + +-- | Apply Householder reflector to one column of a mutable matrix. +-- @rawHouseholderApplyCol mba_r off_r ncols ba_v off_v beta from to col@ +-- computes w = β * Σ_{i=from}^{to-1} v[i] * R[i, col], then +-- R[i, col] -= v[i] * w for i in [from..to-1]. +rawHouseholderApplyCol :: MutableByteArray s -> Int -> Int + -> ByteArray -> Int -> Double + -> Int -> Int -> Int -> ST s () +rawHouseholderApplyCol (MutableByteArray mba_r) (I# off_r) (I# ncols) + (ByteArray ba_v) (I# off_v) (D# beta) + (I# from_) (I# to) (I# col) = ST $ \s0 -> + -- Phase 1: compute w = β * Σ v[i] * R[i, col] + let goSum i acc s + | isTrue# (i >=# to) = (# s, beta *## acc #) + | otherwise = + let vi = indexDoubleArray# ba_v (off_v +# i) + in case readDoubleArray# mba_r (off_r +# i *# ncols +# col) s of + (# s', rij #) -> goSum (i +# 1#) (acc +## vi *## rij) s' + in case goSum from_ 0.0## s0 of + (# s1, w #) -> + -- Phase 2: R[i, col] -= v[i] * w + let goUpdate i s + | isTrue# (i >=# to) = s + | otherwise = + let vi = indexDoubleArray# ba_v (off_v +# i) + in case readDoubleArray# mba_r (off_r +# i *# ncols +# col) s of + (# s', rij #) -> + case writeDoubleArray# mba_r (off_r +# i *# ncols +# col) (rij -## vi *## w) s' of + s'' -> goUpdate (i +# 1#) s'' + in (# goUpdate from_ s1, () #) +{-# INLINE rawHouseholderApplyCol #-} + +-- | Update one row of Q during backward accumulation. +-- @rawQAccumCol mba_q off_q ncols ba_v off_v beta from to row@ +-- computes qi = β * Σ_{k=from}^{to-1} Q[row, k] * v[k], then +-- Q[row, k] -= qi * v[k] for all k in [from..to-1]. +rawQAccumCol :: MutableByteArray s -> Int -> Int + -> ByteArray -> Int -> Double + -> Int -> Int -> Int -> ST s () +rawQAccumCol (MutableByteArray mba_q) (I# off_q) (I# ncols) + (ByteArray ba_v) (I# off_v) (D# beta) + (I# from_) (I# to) (I# row) = ST $ \s0 -> + -- Phase 1: compute qi = β * Σ Q[row, k] * v[k] + let goSum k acc s + | isTrue# (k >=# to) = (# s, beta *## acc #) + | otherwise = + let vk = indexDoubleArray# ba_v (off_v +# k) + in case readDoubleArray# mba_q (off_q +# row *# ncols +# k) s of + (# s', qrk #) -> goSum (k +# 1#) (acc +## vk *## qrk) s' + in case goSum from_ 0.0## s0 of + (# s1, qi #) -> + -- Phase 2: Q[row, k] -= qi * v[k] + let goUpdate k s + | isTrue# (k >=# to) = s + | otherwise = + let vk = indexDoubleArray# ba_v (off_v +# k) + in case readDoubleArray# mba_q (off_q +# row *# ncols +# k) s of + (# s', qrk #) -> + case writeDoubleArray# mba_q (off_q +# row *# ncols +# k) (qrk -## qi *## vk) s' of + s'' -> goUpdate (k +# 1#) s'' + in (# goUpdate from_ s1, () #) +{-# INLINE rawQAccumCol #-} + +-- -------------------------------------------------------------------------- +-- Eigen / tridiag helpers +-- -------------------------------------------------------------------------- + +-- | Apply Givens rotation to two rows of a mutable matrix. +-- @rawApplyGivensRows mba off ncols cosθ sinθ row_p row_q from to@ +-- For each column j in [from..to-1]: +-- tmp = c * M[row_p, j] + s * M[row_q, j] +-- M[row_q, j] = -s * M[row_p, j] + c * M[row_q, j] +-- M[row_p, j] = tmp +rawApplyGivensRows :: MutableByteArray s -> Int -> Int + -> Double -> Double -> Int -> Int + -> Int -> Int -> ST s () +rawApplyGivensRows (MutableByteArray mba) (I# off) (I# ncols) + (D# c_) (D# s_) (I# row_p) (I# row_q) + (I# from_) (I# to) = ST $ \s0 -> + let pOff = off +# row_p *# ncols + qOff = off +# row_q *# ncols + jSpan = to -# from_ + j4End = from_ +# (jSpan -# (jSpan `remInt#` 4#)) + cV = broadcastDoubleX4# c_ + sV = broadcastDoubleX4# s_ + nsV = negateDoubleX4# sV + + goSimd j s + | isTrue# (j >=# j4End) = s + | otherwise = + case readDoubleArrayAsDoubleX4# mba (pOff +# j) s of + (# s1, pv #) -> + case readDoubleArrayAsDoubleX4# mba (qOff +# j) s1 of + (# s2, qv #) -> + let tmp = fmaddDoubleX4# cV pv (timesDoubleX4# sV qv) + q' = fmaddDoubleX4# nsV pv (timesDoubleX4# cV qv) + in case writeDoubleArrayAsDoubleX4# mba (pOff +# j) tmp s2 of + s3 -> case writeDoubleArrayAsDoubleX4# mba (qOff +# j) q' s3 of + s4 -> goSimd (j +# 4#) s4 + + goScalar j s + | isTrue# (j >=# to) = s + | otherwise = + case readDoubleArray# mba (pOff +# j) s of + (# s1, pj #) -> + case readDoubleArray# mba (qOff +# j) s1 of + (# s2, qj #) -> + let tmp = c_ *## pj +## s_ *## qj + qj' = negateDouble# s_ *## pj +## c_ *## qj + in case writeDoubleArray# mba (pOff +# j) tmp s2 of + s3 -> case writeDoubleArray# mba (qOff +# j) qj' s3 of + s4 -> goScalar (j +# 1#) s4 + + in (# goScalar j4End (goSimd from_ s0), () #) +{-# INLINE rawApplyGivensRows #-} + +-- | Symmetric rank-2 update on a mutable matrix. +-- @rawSymRank2Update mba off n ba_v off_v ba_w off_w from to@ +-- For i in [from..to-1], j in [from..i]: +-- T[i,j] -= v[i]*w[j] + w[i]*v[j] +-- T[j,i] = T[i,j] (maintain symmetry) +rawSymRank2Update :: MutableByteArray s -> Int -> Int + -> ByteArray -> Int -> ByteArray -> Int + -> Int -> Int -> ST s () +rawSymRank2Update (MutableByteArray mba) (I# off) (I# n) + (ByteArray ba_v) (I# off_v) (ByteArray ba_w) (I# off_w) + (I# from_) (I# to) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# to) = s + | otherwise = + let vi = indexDoubleArray# ba_v (off_v +# i) + wi = indexDoubleArray# ba_w (off_w +# i) + in goI (i +# 1#) (goJ i vi wi from_ s) + + goJ i vi wi j s + | isTrue# (j ># i) = s + | otherwise = + let vj = indexDoubleArray# ba_v (off_v +# j) + wj = indexDoubleArray# ba_w (off_w +# j) + delta = vi *## wj +## wi *## vj + ij = off +# i *# n +# j + ji = off +# j *# n +# i + in case readDoubleArray# mba ij s of + (# s1, tij #) -> + let tij' = tij -## delta + in case writeDoubleArray# mba ij tij' s1 of + s2 | isTrue# (i ==# j) -> goJ i vi wi (j +# 1#) s2 + | otherwise -> + case writeDoubleArray# mba ji tij' s2 of + s3 -> goJ i vi wi (j +# 1#) s3 + + in (# goI from_ s0, () #) +{-# INLINE rawSymRank2Update #-} + +-- -------------------------------------------------------------------------- +-- Tridiagonalisation mutable kernels +-- -------------------------------------------------------------------------- + +-- | Symmetric submatrix-vector product for tridiagonalisation. +-- Computes p[i-from] = Σ_{j=from}^{to-1} T[i,j] * v[j-from] +-- for i in [from..to-1]. +-- T is read from MutableByteArray (being modified in-place), +-- v is read from MutableByteArray (temporary vector), +-- p is written to MutableByteArray (temporary vector). +rawMutSymMatvecSub :: MutableByteArray s -> Int -> Int + -> MutableByteArray s -> Int + -> MutableByteArray s -> Int + -> Int -> Int -> ST s () +rawMutSymMatvecSub (MutableByteArray mba_t) (I# off_t) (I# ncols) + (MutableByteArray mba_v) (I# off_v) + (MutableByteArray mba_p) (I# off_p) + (I# from_) (I# to) = ST $ \s0 -> + let !len = to -# from_ + !len8 = len -# (len `remInt#` 8#) + !len4 = len -# (len `remInt#` 4#) + + goI i s + | isTrue# (i >=# to) = s + | otherwise = + let !rowBase = off_t +# i *# ncols +# from_ + -- 8-wide SIMD phase: two independent accumulators + in case goJ8 rowBase 0# + (broadcastDoubleX4# 0.0##) + (broadcastDoubleX4# 0.0##) s of + (# s1, accV0, accV1 #) -> + -- 4-wide cleanup phase + case goJ4 rowBase len8 accV0 s1 of + (# s2, accV0' #) -> + -- Reduce both SIMD accumulators to scalar + let !combined = plusDoubleX4# accV0' accV1 + !(# a0, a1, a2, a3 #) = unpackDoubleX4# combined + !simdSum = a0 +## a1 +## a2 +## a3 + -- Scalar tail + in case goJTail rowBase len4 simdSum s2 of + (# s3, acc #) -> + case writeDoubleArray# mba_p (off_p +# (i -# from_)) acc s3 of + s4 -> goI (i +# 1#) s4 + + -- Process 8 doubles (2× DoubleX4#) per iteration + goJ8 rowBase j accV0 accV1 s + | isTrue# (j >=# len8) = (# s, accV0, accV1 #) + | otherwise = + case readDoubleArrayAsDoubleX4# mba_t (rowBase +# j) s of + (# s1, tv0 #) -> + case readDoubleArrayAsDoubleX4# mba_v (off_v +# j) s1 of + (# s2, vv0 #) -> + case readDoubleArrayAsDoubleX4# mba_t (rowBase +# j +# 4#) s2 of + (# s3, tv1 #) -> + case readDoubleArrayAsDoubleX4# mba_v (off_v +# j +# 4#) s3 of + (# s4, vv1 #) -> + goJ8 rowBase (j +# 8#) + (fmaddDoubleX4# tv0 vv0 accV0) + (fmaddDoubleX4# tv1 vv1 accV1) s4 + + -- 4-wide cleanup for elements between len8 and len4 + goJ4 rowBase j accV s + | isTrue# (j >=# len4) = (# s, accV #) + | otherwise = + case readDoubleArrayAsDoubleX4# mba_t (rowBase +# j) s of + (# s1, tv #) -> + case readDoubleArrayAsDoubleX4# mba_v (off_v +# j) s1 of + (# s2, vv #) -> + goJ4 rowBase (j +# 4#) (fmaddDoubleX4# tv vv accV) s2 + + goJTail rowBase j acc s + | isTrue# (j >=# len) = (# s, acc #) + | otherwise = + case readDoubleArray# mba_t (rowBase +# j) s of + (# s1, tij #) -> + case readDoubleArray# mba_v (off_v +# j) s1 of + (# s2, vj #) -> goJTail rowBase (j +# 1#) (acc +## tij *## vj) s2 + + in (# goI from_ s0, () #) +{-# INLINE rawMutSymMatvecSub #-} + +-- | Symmetric rank-2 update reading v, w from MutableByteArrays. +-- For i in [from..to-1], j in [from..i]: +-- T[i,j] -= v[i-from]*w[j-from] + w[i-from]*v[j-from] +-- T[j,i] = T[i,j] (maintain symmetry) +-- v and w are indexed relative to from (i.e., v[0] corresponds to row 'from'). +rawMutSymRank2Update :: MutableByteArray s -> Int -> Int + -> MutableByteArray s -> Int + -> MutableByteArray s -> Int + -> Int -> Int -> ST s () +rawMutSymRank2Update (MutableByteArray mba_t) (I# off_t) (I# n) + (MutableByteArray mba_v) (I# off_v) + (MutableByteArray mba_w) (I# off_w) + (I# from_) (I# to) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# to) = s + | otherwise = + case readDoubleArray# mba_v (off_v +# (i -# from_)) s of + (# s1, vi #) -> + case readDoubleArray# mba_w (off_w +# (i -# from_)) s1 of + (# s2, wi #) -> goI (i +# 1#) (goJ i vi wi from_ s2) + + goJ i vi wi j s + | isTrue# (j ># i) = s + | otherwise = + case readDoubleArray# mba_v (off_v +# (j -# from_)) s of + (# s1, vj #) -> + case readDoubleArray# mba_w (off_w +# (j -# from_)) s1 of + (# s2, wj #) -> + let delta = vi *## wj +## wi *## vj + ij = off_t +# i *# n +# j + ji = off_t +# j *# n +# i + in case readDoubleArray# mba_t ij s2 of + (# s3, tij #) -> + let tij' = tij -## delta + in case writeDoubleArray# mba_t ij tij' s3 of + s4 | isTrue# (i ==# j) -> goJ i vi wi (j +# 1#) s4 + | otherwise -> + case writeDoubleArray# mba_t ji tij' s4 of + s5 -> goJ i vi wi (j +# 1#) s5 + + in (# goI from_ s0, () #) +{-# INLINE rawMutSymRank2Update #-} + +-- | Q accumulation for tridiagonalisation. +-- Householder vectors are stored in column hvCol of frozen T, +-- with implicit v[qCol] = 1.0. +-- Phase 1: wi = beta * (Q[row,qCol] + Σ_{l=qCol+1}^{endRow-1} Q[row,l] * T[l,hvCol]) +-- Phase 2: Q[row,qCol] -= wi; Q[row,l] -= wi * T[l,hvCol] +rawMutTridiagQAccum :: MutableByteArray s -> Int -> Int + -> ByteArray -> Int -> Int + -> Double -> Int -> Int -> Int -> Int -> ST s () +rawMutTridiagQAccum (MutableByteArray mba_q) (I# off_q) (I# qcols) + (ByteArray ba_t) (I# off_t) (I# tcols) + (D# beta) (I# qCol) (I# hvCol) (I# endRow) (I# row) = ST $ \s0 -> + -- Phase 1: wi = beta * (Q[row,qCol] + Σ Q[row,l] * T[l,hvCol]) + case readDoubleArray# mba_q (off_q +# row *# qcols +# qCol) s0 of + (# s1, qrk #) -> + let goSum l acc s + | isTrue# (l >=# endRow) = (# s, beta *## (qrk +## acc) #) + | otherwise = + let vl = indexDoubleArray# ba_t (off_t +# l *# tcols +# hvCol) + in case readDoubleArray# mba_q (off_q +# row *# qcols +# l) s of + (# s', qrl #) -> goSum (l +# 1#) (acc +## qrl *## vl) s' + in case goSum (qCol +# 1#) 0.0## s1 of + (# s2, wi #) -> + -- Phase 2: Q[row,qCol] -= wi + case writeDoubleArray# mba_q (off_q +# row *# qcols +# qCol) (qrk -## wi) s2 of + s3 -> + let goUpdate l s + | isTrue# (l >=# endRow) = s + | otherwise = + let vl = indexDoubleArray# ba_t (off_t +# l *# tcols +# hvCol) + in case readDoubleArray# mba_q (off_q +# row *# qcols +# l) s of + (# s', qrl #) -> + case writeDoubleArray# mba_q (off_q +# row *# qcols +# l) (qrl -## wi *## vl) s' of + s'' -> goUpdate (l +# 1#) s'' + in (# goUpdate (qCol +# 1#) s3, () #) +{-# INLINE rawMutTridiagQAccum #-} + +-- -------------------------------------------------------------------------- +-- LU kernels +-- -------------------------------------------------------------------------- + +-- | In-place LU elimination for column k of an n×n row-major matrix. +-- Computes multipliers and updates the trailing submatrix. +-- Inner j-loop uses DoubleX4# SIMD (contiguous row access). +rawLUEliminateColumn :: MutableByteArray s -> Int -> Int -> Int -> ST s () +rawLUEliminateColumn (MutableByteArray mba) (I# off) (I# n) (I# k) = ST $ \s0 -> + -- Read A[k,k] (the pivot) + case readDoubleArray# mba (off +# k *# n +# k) s0 of + (# s1, akk #) -> + let goI i s + | isTrue# (i >=# n) = s + | otherwise = + -- Read A[i,k], compute multiplier + case readDoubleArray# mba (off +# i *# n +# k) s of + (# s', aik #) -> + let mult = aik /## akk + iRowOff = off +# i *# n + kRowOff = off +# k *# n + jSpan = n -# k -# 1# + jStart = k +# 1# + j4End = jStart +# (jSpan -# (jSpan `remInt#` 4#)) + negMultV = broadcastDoubleX4# (negateDouble# mult) + -- Store multiplier at A[i,k] + in case writeDoubleArray# mba (off +# i *# n +# k) mult s' of + s'' -> + -- SIMD j-loop: A[i,j] -= mult * A[k,j] = A[i,j] + (-mult)*A[k,j] + let goJSimd j s_ + | isTrue# (j >=# j4End) = s_ + | otherwise = + case readDoubleArrayAsDoubleX4# mba (iRowOff +# j) s_ of + (# s1_, aij #) -> + case readDoubleArrayAsDoubleX4# mba (kRowOff +# j) s1_ of + (# s2_, akjV_ #) -> + let aij' = fmaddDoubleX4# negMultV akjV_ aij + in case writeDoubleArrayAsDoubleX4# mba (iRowOff +# j) aij' s2_ of + s3_ -> goJSimd (j +# 4#) s3_ + -- Scalar cleanup + goJScalar j s_ + | isTrue# (j >=# n) = s_ + | otherwise = + case readDoubleArray# mba (iRowOff +# j) s_ of + (# s1_, aij #) -> + case readDoubleArray# mba (kRowOff +# j) s1_ of + (# s2_, akj #) -> + case writeDoubleArray# mba (iRowOff +# j) (aij -## mult *## akj) s2_ of + s3_ -> goJScalar (j +# 1#) s3_ + in goI (i +# 1#) (goJScalar j4End (goJSimd jStart s'')) + in (# goI (k +# 1#) s1, () #) +{-# INLINE rawLUEliminateColumn #-} + +-- | @rawLUEliminateColumnTo mba off n k colEnd@ — like 'rawLUEliminateColumn' +-- but the trailing update only touches columns @k+1 .. colEnd-1@ (not @k+1 .. n-1@). +-- Multipliers are still computed for ALL rows @k+1 .. n-1@. +-- Used by panel LU to restrict updates to within the current panel. +rawLUEliminateColumnTo :: MutableByteArray s -> Int -> Int -> Int -> Int -> ST s () +rawLUEliminateColumnTo (MutableByteArray mba) (I# off) (I# n) (I# k) (I# colEnd) = ST $ \s0 -> + case readDoubleArray# mba (off +# k *# n +# k) s0 of + (# s1, akk #) -> + let goI i s + | isTrue# (i >=# n) = s + | otherwise = + case readDoubleArray# mba (off +# i *# n +# k) s of + (# s', aik #) -> + let mult = aik /## akk + iRowOff = off +# i *# n + kRowOff = off +# k *# n + jSpan = colEnd -# k -# 1# + jStart = k +# 1# + j4End = jStart +# (jSpan -# (jSpan `remInt#` 4#)) + negMultV = broadcastDoubleX4# (negateDouble# mult) + in case writeDoubleArray# mba (off +# i *# n +# k) mult s' of + s'' -> + let goJSimd j s_ + | isTrue# (j >=# j4End) = s_ + | otherwise = + case readDoubleArrayAsDoubleX4# mba (iRowOff +# j) s_ of + (# s1_, aij #) -> + case readDoubleArrayAsDoubleX4# mba (kRowOff +# j) s1_ of + (# s2_, akjV_ #) -> + let aij' = fmaddDoubleX4# negMultV akjV_ aij + in case writeDoubleArrayAsDoubleX4# mba (iRowOff +# j) aij' s2_ of + s3_ -> goJSimd (j +# 4#) s3_ + goJScalar j s_ + | isTrue# (j >=# colEnd) = s_ + | otherwise = + case readDoubleArray# mba (iRowOff +# j) s_ of + (# s1_, aij #) -> + case readDoubleArray# mba (kRowOff +# j) s1_ of + (# s2_, akj #) -> + case writeDoubleArray# mba (iRowOff +# j) (aij -## mult *## akj) s2_ of + s3_ -> goJScalar (j +# 1#) s3_ + in goI (i +# 1#) (goJScalar j4End (goJSimd jStart s'')) + in (# goI (k +# 1#) s1, () #) +{-# INLINE rawLUEliminateColumnTo #-} + +-- | Swap elements in columns [fromCol..n-1] between two rows of an n-wide matrix. +-- Uses DoubleX4# SIMD for the contiguous row data. +rawSwapRows :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> ST s () +rawSwapRows (MutableByteArray mba) (I# off) (I# n) (I# row1) (I# row2) (I# fromCol) = ST $ \s0 -> + let r1Off = off +# row1 *# n + r2Off = off +# row2 *# n + jSpan = n -# fromCol + j4End = fromCol +# (jSpan -# (jSpan `remInt#` 4#)) + + goSimd j s + | isTrue# (j >=# j4End) = s + | otherwise = + case readDoubleArrayAsDoubleX4# mba (r1Off +# j) s of + (# s1, v1 #) -> + case readDoubleArrayAsDoubleX4# mba (r2Off +# j) s1 of + (# s2, v2 #) -> + case writeDoubleArrayAsDoubleX4# mba (r1Off +# j) v2 s2 of + s3 -> case writeDoubleArrayAsDoubleX4# mba (r2Off +# j) v1 s3 of + s4 -> goSimd (j +# 4#) s4 + + goScalar j s + | isTrue# (j >=# n) = s + | otherwise = + case readDoubleArray# mba (r1Off +# j) s of + (# s1, v1 #) -> + case readDoubleArray# mba (r2Off +# j) s1 of + (# s2, v2 #) -> + case writeDoubleArray# mba (r1Off +# j) v2 s2 of + s3 -> case writeDoubleArray# mba (r2Off +# j) v1 s3 of + s4 -> goScalar (j +# 1#) s4 + + in (# goScalar j4End (goSimd fromCol s0), () #) +{-# INLINE rawSwapRows #-} + +-- | Find row with maximum |A[i,k]| for i in [fromRow..n-1]. +-- Returns the row index of the pivot. +rawPivotSearch :: MutableByteArray s -> Int -> Int -> Int -> Int -> ST s Int +rawPivotSearch (MutableByteArray mba) (I# off) (I# n) (I# k) (I# fromRow) = ST $ \s0 -> + let go i bestIdx bestVal s + | isTrue# (i >=# n) = (# s, I# bestIdx #) + | otherwise = + case readDoubleArray# mba (off +# i *# n +# k) s of + (# s', v #) -> + let av = if isTrue# (v >=## 0.0##) then v else negateDouble# v + in if isTrue# (av >## bestVal) + then go (i +# 1#) i av s' + else go (i +# 1#) bestIdx bestVal s' + in case readDoubleArray# mba (off +# fromRow *# n +# k) s0 of + (# s1, v0 #) -> + let av0 = if isTrue# (v0 >=## 0.0##) then v0 else negateDouble# v0 + in go (fromRow +# 1#) fromRow av0 s1 +{-# INLINE rawPivotSearch #-} + +-- | In-place forward substitution using packed LU matrix (unit lower triangular). +-- Solves Ly = b where L is stored in the strictly lower part of ba_lu. +-- The solution overwrites mba_x. +rawForwardSubUnitPacked :: ByteArray -> Int -> Int -> MutableByteArray s -> Int -> ST s () +rawForwardSubUnitPacked (ByteArray ba_lu) (I# off_lu) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goJ j s + | isTrue# (j >=# n) = s + | otherwise = + case readDoubleArray# mba_x (off_x +# j) s of + (# s', xj #) -> + let goI i s_ + | isTrue# (i >=# n) = s_ + | otherwise = + let lij = indexDoubleArray# ba_lu (off_lu +# i *# n +# j) + in case readDoubleArray# mba_x (off_x +# i) s_ of + (# s1, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) (xi -## lij *## xj) s1 of + s2 -> goI (i +# 1#) s2 + in goJ (j +# 1#) (goI (j +# 1#) s') + in (# goJ 0# s0, () #) +{-# INLINE rawForwardSubUnitPacked #-} + +-- | In-place back substitution using packed LU matrix (upper triangular). +-- Solves Ux = y where U is stored in the upper part of ba_lu. +-- The solution overwrites mba_x. +rawBackSubPacked :: ByteArray -> Int -> Int -> MutableByteArray s -> Int -> ST s () +rawBackSubPacked (ByteArray ba_lu) (I# off_lu) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goJ j s + | isTrue# (j <# 0#) = s + | otherwise = + -- x[j] /= U[j,j] + let ujj = indexDoubleArray# ba_lu (off_lu +# j *# n +# j) + in case readDoubleArray# mba_x (off_x +# j) s of + (# s', xj_ #) -> + let xj = xj_ /## ujj + in case writeDoubleArray# mba_x (off_x +# j) xj s' of + s'' -> + -- for i = 0..j-1: x[i] -= U[i,j] * x[j] + let goI i s_ + | isTrue# (i >=# j) = s_ + | otherwise = + let uij = indexDoubleArray# ba_lu (off_lu +# i *# n +# j) + in case readDoubleArray# mba_x (off_x +# i) s_ of + (# s1, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) (xi -## uij *## xj) s1 of + s2 -> goI (i +# 1#) s2 + in goJ (j -# 1#) (goI 0# s'') + in (# goJ (n -# 1#) s0, () #) +{-# INLINE rawBackSubPacked #-} + +-- -------------------------------------------------------------------------- +-- Cholesky kernels +-- -------------------------------------------------------------------------- + +-- | Process one column j of Cholesky factorisation in-place. +-- For k in [0..j-1]: subtract G[i,k]*G[j,k] from G[i,j] for i in [j..n-1]. +-- Then scale: G[j,j] = sqrt(G[j,j]); G[i,j] /= G[j,j] for i > j. +rawCholColumn :: MutableByteArray s -> Int -> Int -> Int -> ST s () +rawCholColumn (MutableByteArray mba) (I# off) (I# n) (I# j) = ST $ \s0 -> + -- Phase 1: subtract contributions from previous columns + let goK k s + | isTrue# (k >=# j) = s + | otherwise = + -- Read G[j,k] + case readDoubleArray# mba (off +# j *# n +# k) s of + (# s', gjk #) -> + -- For i in [j..n-1]: G[i,j] -= G[i,k] * gjk + let goI i s_ + | isTrue# (i >=# n) = s_ + | otherwise = + case readDoubleArray# mba (off +# i *# n +# j) s_ of + (# s1, gij #) -> + case readDoubleArray# mba (off +# i *# n +# k) s1 of + (# s2, gik #) -> + case writeDoubleArray# mba (off +# i *# n +# j) (gij -## gik *## gjk) s2 of + s3 -> goI (i +# 1#) s3 + in goK (k +# 1#) (goI j s') + in case goK 0# s0 of + s1 -> + -- Phase 2: scale column + case readDoubleArray# mba (off +# j *# n +# j) s1 of + (# s2, gjj #) -> + let sjj = sqrtDouble# gjj + in case writeDoubleArray# mba (off +# j *# n +# j) sjj s2 of + s3 -> + let goScale i s_ + | isTrue# (i >=# n) = s_ + | otherwise = + case readDoubleArray# mba (off +# i *# n +# j) s_ of + (# s4, gij #) -> + case writeDoubleArray# mba (off +# i *# n +# j) (gij /## sjj) s4 of + s5 -> goScale (i +# 1#) s5 + in (# goScale (j +# 1#) s3, () #) +{-# INLINE rawCholColumn #-} + +-- | Forward substitution with Cholesky factor G (lower triangular, non-unit diagonal). +-- Solves Gy = b, overwrites mba_x with y. +rawForwardSubCholPacked :: ByteArray -> Int -> Int -> MutableByteArray s -> Int -> ST s () +rawForwardSubCholPacked (ByteArray ba_g) (I# off_g) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goJ j s + | isTrue# (j >=# n) = s + | otherwise = + let gjj = indexDoubleArray# ba_g (off_g +# j *# n +# j) + in case readDoubleArray# mba_x (off_x +# j) s of + (# s', xj_ #) -> + let xj = xj_ /## gjj + in case writeDoubleArray# mba_x (off_x +# j) xj s' of + s'' -> + let goI i s_ + | isTrue# (i >=# n) = s_ + | otherwise = + let gij = indexDoubleArray# ba_g (off_g +# i *# n +# j) + in case readDoubleArray# mba_x (off_x +# i) s_ of + (# s1, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) (xi -## gij *## xj) s1 of + s2 -> goI (i +# 1#) s2 + in goJ (j +# 1#) (goI (j +# 1#) s'') + in (# goJ 0# s0, () #) +{-# INLINE rawForwardSubCholPacked #-} + +-- | Back substitution with G^T (upper triangular) WITHOUT forming G^T. +-- Solves G^T x = y, overwrites mba_x with x. +-- Uses G^T[i,j] = G[j,i] to read from the lower triangle. +rawBackSubCholTPacked :: ByteArray -> Int -> Int -> MutableByteArray s -> Int -> ST s () +rawBackSubCholTPacked (ByteArray ba_g) (I# off_g) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goJ j s + | isTrue# (j <# 0#) = s + | otherwise = + -- G^T[j,j] = G[j,j] + let gjj = indexDoubleArray# ba_g (off_g +# j *# n +# j) + in case readDoubleArray# mba_x (off_x +# j) s of + (# s', xj_ #) -> + let xj = xj_ /## gjj + in case writeDoubleArray# mba_x (off_x +# j) xj s' of + s'' -> + -- for i = 0..j-1: x[i] -= G^T[i,j] * x[j] = G[j,i] * x[j] + let goI i s_ + | isTrue# (i >=# j) = s_ + | otherwise = + -- G^T[i,j] = G[j,i] + let gji = indexDoubleArray# ba_g (off_g +# j *# n +# i) + in case readDoubleArray# mba_x (off_x +# i) s_ of + (# s1, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) (xi -## gji *## xj) s1 of + s2 -> goI (i +# 1#) s2 + in goJ (j -# 1#) (goI 0# s'') + in (# goJ (n -# 1#) s0, () #) +{-# INLINE rawBackSubCholTPacked #-} + +-- -------------------------------------------------------------------------- +-- QR mutable kernels +-- -------------------------------------------------------------------------- + +-- | Sum of squares of a column slice in a mutable row-major matrix. +-- Σ A[i,col]² for i in [startRow..endRow-1]. +rawMutSumSqColumn :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> ST s Double +rawMutSumSqColumn (MutableByteArray mba) (I# off) (I# ncols) (I# startRow) (I# endRow) (I# col) = ST $ \s0 -> + let go i acc s + | isTrue# (i >=# endRow) = (# s, D# acc #) + | otherwise = + case readDoubleArray# mba (off +# i *# ncols +# col) s of + (# s', v #) -> go (i +# 1#) (acc +## v *## v) s' + in go startRow 0.0## s0 +{-# INLINE rawMutSumSqColumn #-} + +-- | Dot product of two column slices in a mutable row-major matrix. +-- Σ A[i,col1] * A[i,col2] for i in [startRow..endRow-1]. +rawMutSumProdColumns :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> Int -> ST s Double +rawMutSumProdColumns (MutableByteArray mba) (I# off) (I# ncols) (I# startRow) (I# endRow) (I# col1) (I# col2) = ST $ \s0 -> + let go i acc s + | isTrue# (i >=# endRow) = (# s, D# acc #) + | otherwise = + case readDoubleArray# mba (off +# i *# ncols +# col1) s of + (# s1, v1 #) -> + case readDoubleArray# mba (off +# i *# ncols +# col2) s1 of + (# s2, v2 #) -> go (i +# 1#) (acc +## v1 *## v2) s2 + in go startRow 0.0## s0 +{-# INLINE rawMutSumProdColumns #-} + +-- | Apply Householder reflector stored in column k (rows k+1..endRow-1, +-- with v[k]=1 implicit) to targetCol of a mutable row-major matrix. +-- Phase 1: w = beta * (R[k,targetCol] + Σ_{i=k+1}^{endRow-1} v[i]*R[i,targetCol]) +-- Phase 2: R[k,targetCol] -= w; R[i,targetCol] -= v[i]*w +rawMutHouseholderApply :: MutableByteArray s -> Int -> Int -> Double + -> Int -> Int -> Int -> ST s () +rawMutHouseholderApply (MutableByteArray mba) (I# off) (I# ncols) (D# beta) + (I# k) (I# endRow) (I# targetCol) = ST $ \s0 -> + -- Phase 1: compute dot product + case readDoubleArray# mba (off +# k *# ncols +# targetCol) s0 of + (# s1, rkj #) -> + let goSum i acc s + | isTrue# (i >=# endRow) = (# s, beta *## (rkj +## acc) #) + | otherwise = + case readDoubleArray# mba (off +# i *# ncols +# k) s of + (# s', vi #) -> + case readDoubleArray# mba (off +# i *# ncols +# targetCol) s' of + (# s'', rij #) -> goSum (i +# 1#) (acc +## vi *## rij) s'' + in case goSum (k +# 1#) 0.0## s1 of + (# s2, w #) -> + -- Phase 2: update R[k,targetCol] + case writeDoubleArray# mba (off +# k *# ncols +# targetCol) (rkj -## w) s2 of + s3 -> + let goUpdate i s + | isTrue# (i >=# endRow) = s + | otherwise = + case readDoubleArray# mba (off +# i *# ncols +# k) s of + (# s', vi #) -> + case readDoubleArray# mba (off +# i *# ncols +# targetCol) s' of + (# s'', rij #) -> + case writeDoubleArray# mba (off +# i *# ncols +# targetCol) (rij -## vi *## w) s'' of + s''' -> goUpdate (i +# 1#) s''' + in (# goUpdate (k +# 1#) s3, () #) +{-# INLINE rawMutHouseholderApply #-} + +-- | Apply stored Householder reflector to one row of Q during accumulation. +-- v is stored in the subdiagonal of frozen R (column k, rows k+1..endRow-1). +-- Phase 1: wi = beta * (Q[row,k] + Σ_{l=k+1}^{endRow-1} Q[row,l] * v[l]) +-- Phase 2: Q[row,k] -= wi; Q[row,l] -= wi * v[l] +rawMutQAccum :: MutableByteArray s -> Int -> Int + -> ByteArray -> Int -> Int + -> Double -> Int -> Int -> Int -> ST s () +rawMutQAccum (MutableByteArray mba_q) (I# off_q) (I# qcols) + (ByteArray ba_r) (I# off_r) (I# rcols) + (D# beta) (I# k) (I# endRow) (I# row) = ST $ \s0 -> + -- Phase 1: compute wi = beta * (Q[row,k] + Σ Q[row,l] * v[l]) + case readDoubleArray# mba_q (off_q +# row *# qcols +# k) s0 of + (# s1, qrk #) -> + let goSum l acc s + | isTrue# (l >=# endRow) = (# s, beta *## (qrk +## acc) #) + | otherwise = + let vl = indexDoubleArray# ba_r (off_r +# l *# rcols +# k) + in case readDoubleArray# mba_q (off_q +# row *# qcols +# l) s of + (# s', qrl #) -> goSum (l +# 1#) (acc +## qrl *## vl) s' + in case goSum (k +# 1#) 0.0## s1 of + (# s2, wi #) -> + -- Phase 2: Q[row,k] -= wi + case writeDoubleArray# mba_q (off_q +# row *# qcols +# k) (qrk -## wi) s2 of + s3 -> + let goUpdate l s + | isTrue# (l >=# endRow) = s + | otherwise = + let vl = indexDoubleArray# ba_r (off_r +# l *# rcols +# k) + in case readDoubleArray# mba_q (off_q +# row *# qcols +# l) s of + (# s', qrl #) -> + case writeDoubleArray# mba_q (off_q +# row *# qcols +# l) (qrl -## wi *## vl) s' of + s'' -> goUpdate (l +# 1#) s'' + in (# goUpdate (k +# 1#) s3, () #) +{-# INLINE rawMutQAccum #-} + +-- -------------------------------------------------------------------------- +-- Eigen mutable kernels +-- -------------------------------------------------------------------------- + +-- | Apply Givens rotation to two columns of a mutable matrix. +-- For each row in [0..nrows-1]: +-- tmp = c * M[row,col_p] + s * M[row,col_q] +-- M[row,col_q] = -s * M[row,col_p] + c * M[row,col_q] +-- M[row,col_p] = tmp +rawMutApplyGivensColumns :: MutableByteArray s -> Int -> Int + -> Double -> Double -> Int -> Int -> Int -> ST s () +rawMutApplyGivensColumns (MutableByteArray mba) (I# off) (I# ncols) + (D# c_) (D# s_) (I# col_p) (I# col_q) (I# nrows) = ST $ \s0 -> + let go row s + | isTrue# (row >=# nrows) = s + | otherwise = + let pIdx = off +# row *# ncols +# col_p + qIdx = off +# row *# ncols +# col_q + in case readDoubleArray# mba pIdx s of + (# s1, mp #) -> + case readDoubleArray# mba qIdx s1 of + (# s2, mq #) -> + let tmp = c_ *## mp +## s_ *## mq + qnew = negateDouble# s_ *## mp +## c_ *## mq + in case writeDoubleArray# mba pIdx tmp s2 of + s3 -> case writeDoubleArray# mba qIdx qnew s3 of + s4 -> go (row +# 1#) s4 + in (# go 0# s0, () #) +{-# INLINE rawMutApplyGivensColumns #-} + +-- | Apply Givens rotation to two columns of a COLUMN-MAJOR mutable matrix. +-- In column-major layout, Q[i,j] is at off + j*nrows + i. +-- Column col_p occupies contiguous memory, enabling SIMD vectorisation. +-- For each row in [0..nrows-1]: +-- tmp = c * M[row,col_p] + s * M[row,col_q] +-- M[row,col_q] = -s * M[row,col_p] + c * M[row,col_q] +-- M[row,col_p] = tmp +rawMutApplyGivensColumnsCM :: MutableByteArray s -> Int -> Int + -> Double -> Double -> Int -> Int -> Int -> ST s () +rawMutApplyGivensColumnsCM (MutableByteArray mba) (I# off) (I# nrows) + (D# c_) (D# s_) (I# col_p) (I# col_q) (I# _ncols) = ST $ \s0 -> + let pBase = off +# col_p *# nrows + qBase = off +# col_q *# nrows + nrows4 = nrows -# (nrows `remInt#` 4#) + cV = broadcastDoubleX4# c_ + sV = broadcastDoubleX4# s_ + nsV = negateDoubleX4# sV + + goSimd i s + | isTrue# (i >=# nrows4) = s + | otherwise = + case readDoubleArrayAsDoubleX4# mba (pBase +# i) s of + (# s1, pv #) -> + case readDoubleArrayAsDoubleX4# mba (qBase +# i) s1 of + (# s2, qv #) -> + let tmp = fmaddDoubleX4# cV pv (timesDoubleX4# sV qv) + q' = fmaddDoubleX4# nsV pv (timesDoubleX4# cV qv) + in case writeDoubleArrayAsDoubleX4# mba (pBase +# i) tmp s2 of + s3 -> case writeDoubleArrayAsDoubleX4# mba (qBase +# i) q' s3 of + s4 -> goSimd (i +# 4#) s4 + + goScalar i s + | isTrue# (i >=# nrows) = s + | otherwise = + case readDoubleArray# mba (pBase +# i) s of + (# s1, mp #) -> + case readDoubleArray# mba (qBase +# i) s1 of + (# s2, mq #) -> + let tmp = c_ *## mp +## s_ *## mq + qnew = negateDouble# s_ *## mp +## c_ *## mq + in case writeDoubleArray# mba (pBase +# i) tmp s2 of + s3 -> case writeDoubleArray# mba (qBase +# i) qnew s3 of + s4 -> goScalar (i +# 1#) s4 + + in (# goScalar nrows4 (goSimd 0# s0), () #) +{-# INLINE rawMutApplyGivensColumnsCM #-} + +-- -------------------------------------------------------------------------- +-- Matrix transpose (row-major <-> column-major) +-- -------------------------------------------------------------------------- + +-- | Transpose an n×n row-major matrix to column-major layout. +-- src[i,j] at offS + i*n + j -> dst[i,j] at offD + j*n + i +rawTransposeToColMajor :: MutableByteArray s -> Int -> MutableByteArray s -> Int -> Int -> ST s () +rawTransposeToColMajor (MutableByteArray src) (I# offS) + (MutableByteArray dst) (I# offD) (I# n) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# n) = s + | otherwise = + let goJ j s' + | isTrue# (j >=# n) = s' + | otherwise = + case readDoubleArray# src (offS +# i *# n +# j) s' of + (# s1, v #) -> + case writeDoubleArray# dst (offD +# j *# n +# i) v s1 of + s2 -> goJ (j +# 1#) s2 + in goI (i +# 1#) (goJ 0# s) + in (# goI 0# s0, () #) +{-# INLINE rawTransposeToColMajor #-} + +-- | Transpose an n×n column-major matrix back to row-major layout. +-- src[i,j] at offS + j*n + i -> dst[i,j] at offD + i*n + j +rawTransposeFromColMajor :: MutableByteArray s -> Int -> MutableByteArray s -> Int -> Int -> ST s () +rawTransposeFromColMajor (MutableByteArray src) (I# offS) + (MutableByteArray dst) (I# offD) (I# n) = ST $ \s0 -> + let goJ j s + | isTrue# (j >=# n) = s + | otherwise = + let goI i s' + | isTrue# (i >=# n) = s' + | otherwise = + case readDoubleArray# src (offS +# j *# n +# i) s' of + (# s1, v #) -> + case writeDoubleArray# dst (offD +# i *# n +# j) v s1 of + s2 -> goI (i +# 1#) s2 + in goJ (j +# 1#) (goI 0# s) + in (# goJ 0# s0, () #) +{-# INLINE rawTransposeFromColMajor #-} + +-- -------------------------------------------------------------------------- +-- Bulk memory operations +-- -------------------------------------------------------------------------- + +-- | Zero n consecutive doubles in a MutableByteArray starting at element offset. +-- Uses SIMD (DoubleX4#) for the main loop with scalar cleanup. +rawZeroDoubles :: MutableByteArray s -> Int -> Int -> ST s () +rawZeroDoubles (MutableByteArray mba) (I# off) (I# n) = ST $ \s0 -> + let n4 = n -# (n `remInt#` 4#) + zeroV = broadcastDoubleX4# 0.0## + + goSimd i s + | isTrue# (i >=# n4) = s + | otherwise = + case writeDoubleArrayAsDoubleX4# mba (off +# i) zeroV s of + s1 -> goSimd (i +# 4#) s1 + + goScalar i s + | isTrue# (i >=# n) = s + | otherwise = + case writeDoubleArray# mba (off +# i) 0.0## s of + s1 -> goScalar (i +# 1#) s1 + + in (# goScalar n4 (goSimd 0# s0), () #) +{-# INLINE rawZeroDoubles #-} + +-- | Copy n consecutive doubles from src to dst using memcpy (copyMutableByteArray#). +-- @rawCopyDoubles dst dstOff src srcOff n@ copies src[srcOff..srcOff+n-1] to dst[dstOff..dstOff+n-1]. +-- All offsets are in element (Double) units. +rawCopyDoubles :: MutableByteArray s -> Int -> MutableByteArray s -> Int -> Int -> ST s () +rawCopyDoubles (MutableByteArray dst) (I# dstOff) (MutableByteArray src) (I# srcOff) (I# n) = ST $ \s -> + case copyMutableByteArray# src (srcOff *# 8#) dst (dstOff *# 8#) (n *# 8#) s of + s' -> (# s', () #) +{-# INLINE rawCopyDoubles #-} + +-- | Negate n consecutive doubles in-place using SIMD. +rawNegateDoubles :: MutableByteArray s -> Int -> Int -> ST s () +rawNegateDoubles (MutableByteArray mba) (I# off) (I# n) = ST $ \s0 -> + let n4 = n -# (n `remInt#` 4#) + negOneV = broadcastDoubleX4# (negateDouble# 1.0##) + + goSimd i s + | isTrue# (i >=# n4) = s + | otherwise = + case readDoubleArrayAsDoubleX4# mba (off +# i) s of + (# s1, v #) -> + case writeDoubleArrayAsDoubleX4# mba (off +# i) (timesDoubleX4# negOneV v) s1 of + s2 -> goSimd (i +# 4#) s2 + + goScalar i s + | isTrue# (i >=# n) = s + | otherwise = + case readDoubleArray# mba (off +# i) s of + (# s1, v #) -> + case writeDoubleArray# mba (off +# i) (negateDouble# v) s1 of + s2 -> goScalar (i +# 1#) s2 + + in (# goScalar n4 (goSimd 0# s0), () #) +{-# INLINE rawNegateDoubles #-} + +-- | Copy a column from one matrix to another (both row-major). +-- Copies src[row, srcCol] to dst[row, dstCol] for row in [0..nrows-1]. +-- Parameters: srcMBA srcOff srcStride srcCol -> dstMBA dstOff dstStride dstCol -> nrows +rawCopyColumn :: MutableByteArray s -> Int -> Int -> Int + -> MutableByteArray s -> Int -> Int -> Int -> Int -> ST s () +rawCopyColumn (MutableByteArray src) (I# offS) (I# strideS) (I# colS) + (MutableByteArray dst) (I# offD) (I# strideD) (I# colD) (I# nrows) = ST $ \s0 -> + let go i s + | isTrue# (i >=# nrows) = s + | otherwise = + case readDoubleArray# src (offS +# i *# strideS +# colS) s of + (# s1, v #) -> + case writeDoubleArray# dst (offD +# i *# strideD +# colD) v s1 of + s2 -> go (i +# 1#) s2 + in (# go 0# s0, () #) +{-# INLINE rawCopyColumn #-} + +-- -------------------------------------------------------------------------- +-- SVD / bidiagonalisation kernels +-- -------------------------------------------------------------------------- + +-- | Apply a right Householder reflector to one row of a mutable matrix. +-- The Householder vector v is stored in row hvRow of the matrix, +-- columns [hvStart..hvEnd-1], with implicit v[hvStart] = 1.0. +-- Updates row targetRow: R[targetRow, hvStart..hvEnd-1] -= w * v +-- where w = beta * (R[targetRow,hvStart] + Σ_{l=hvStart+1}^{hvEnd-1} R[targetRow,l] * R[hvRow,l]) +rawMutHouseholderApplyRow :: MutableByteArray s -> Int -> Int + -> Double -> Int -> Int -> Int -> Int -> ST s () +rawMutHouseholderApplyRow (MutableByteArray mba) (I# off) (I# ncols) (D# beta) + (I# hvRow) (I# hvStart) (I# hvEnd) (I# targetRow) = ST $ \s0 -> + let trOff = off +# targetRow *# ncols + hvOff = off +# hvRow *# ncols + -- Phase 1: w = beta * (R[targetRow,hvStart] + Σ R[targetRow,l] * R[hvRow,l]) + in case readDoubleArray# mba (trOff +# hvStart) s0 of + (# s1, r0 #) -> + let goSum l acc s + | isTrue# (l >=# hvEnd) = (# s, beta *## (r0 +## acc) #) + | otherwise = + case readDoubleArray# mba (trOff +# l) s of + (# s', rl #) -> + case readDoubleArray# mba (hvOff +# l) s' of + (# s'', vl #) -> goSum (l +# 1#) (acc +## rl *## vl) s'' + in case goSum (hvStart +# 1#) 0.0## s1 of + (# s2, w #) -> + -- Phase 2: R[targetRow,hvStart] -= w (implicit v[hvStart]=1) + case writeDoubleArray# mba (trOff +# hvStart) (r0 -## w) s2 of + s3 -> + let goUpdate l s + | isTrue# (l >=# hvEnd) = s + | otherwise = + case readDoubleArray# mba (hvOff +# l) s of + (# s', vl #) -> + case readDoubleArray# mba (trOff +# l) s' of + (# s'', rl #) -> + case writeDoubleArray# mba (trOff +# l) (rl -## w *## vl) s'' of + s''' -> goUpdate (l +# 1#) s''' + in (# goUpdate (hvStart +# 1#) s3, () #) +{-# INLINE rawMutHouseholderApplyRow #-} + +-- | Sum of squares of a row slice in a mutable row-major matrix. +-- Σ A[row,j]² for j in [startCol..endCol-1]. +rawMutSumSqRow :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> ST s Double +rawMutSumSqRow (MutableByteArray mba) (I# off) (I# ncols) (I# row) (I# startCol) (I# endCol) = ST $ \s0 -> + let rowOff = off +# row *# ncols + go j acc s + | isTrue# (j >=# endCol) = (# s, D# acc #) + | otherwise = + case readDoubleArray# mba (rowOff +# j) s of + (# s', v #) -> go (j +# 1#) (acc +## v *## v) s' + in go startCol 0.0## s0 +{-# INLINE rawMutSumSqRow #-} + +-- -------------------------------------------------------------------------- +-- Cholesky SIMD kernel +-- -------------------------------------------------------------------------- + +-- | SIMD-vectorised Cholesky column kernel. +-- Restructures the inner loop as a dot product of contiguous row segments: +-- G[i,j] -= Σ_{k=0}^{j-1} G[i,k] * G[j,k] +-- which is a dot product of row[i][0..j-1] and row[j][0..j-1]. +-- Since rows are contiguous in row-major storage, this enables DoubleX4# SIMD. +rawCholColumnSIMD :: MutableByteArray s -> Int -> Int -> Int -> ST s () +rawCholColumnSIMD (MutableByteArray mba) (I# off) (I# n) (I# j) = ST $ \s0 -> + let jRowOff = off +# j *# n + -- For each row i in [j..n-1], subtract dot(row[i][0..j-1], row[j][0..j-1]) + j4 = j -# (j `remInt#` 4#) -- SIMD boundary for dot of length j + + goI i s + | isTrue# (i >=# n) = s + | otherwise = + let iRowOff = off +# i *# n + in case mutRowDot iRowOff jRowOff 0# j4 j s of + (# s', dot #) -> + case readDoubleArray# mba (iRowOff +# j) s' of + (# s'', gij #) -> + case writeDoubleArray# mba (iRowOff +# j) (gij -## dot) s'' of + s''' -> goI (i +# 1#) s''' + + -- Dot product of two mutable row segments using SIMD + mutRowDot r1 r2 k k4End kEnd s + -- SIMD phase + | isTrue# (k <# k4End) = + case goSimd r1 r2 k k4End (broadcastDoubleX4# 0.0##) s of + (# s', acc4 #) -> + let !(# a, b, c, d #) = unpackDoubleX4# acc4 + simdSum = a +## b +## c +## d + in mutRowDotScalar r1 r2 k4End kEnd simdSum s' + | otherwise = mutRowDotScalar r1 r2 k kEnd 0.0## s + + goSimd r1 r2 k k4End acc s + | isTrue# (k >=# k4End) = (# s, acc #) + | otherwise = + case readDoubleArrayAsDoubleX4# mba (r1 +# k) s of + (# s1, v1 #) -> + case readDoubleArrayAsDoubleX4# mba (r2 +# k) s1 of + (# s2, v2 #) -> goSimd r1 r2 (k +# 4#) k4End (fmaddDoubleX4# v1 v2 acc) s2 + + mutRowDotScalar r1 r2 k kEnd acc s + | isTrue# (k >=# kEnd) = (# s, acc #) + | otherwise = + case readDoubleArray# mba (r1 +# k) s of + (# s1, v1 #) -> + case readDoubleArray# mba (r2 +# k) s1 of + (# s2, v2 #) -> mutRowDotScalar r1 r2 (k +# 1#) kEnd (acc +## v1 *## v2) s2 + + in case goI j s0 of + s1 -> + -- Scale column: G[j,j] = sqrt(G[j,j]); G[i,j] /= G[j,j] for i > j + case readDoubleArray# mba (jRowOff +# j) s1 of + (# s2, gjj #) -> + let sjj = sqrtDouble# gjj + in case writeDoubleArray# mba (jRowOff +# j) sjj s2 of + s3 -> + let goScale i s + | isTrue# (i >=# n) = s + | otherwise = + case readDoubleArray# mba (off +# i *# n +# j) s of + (# s4, gij #) -> + case writeDoubleArray# mba (off +# i *# n +# j) (gij /## sjj) s4 of + s5 -> goScale (i +# 1#) s5 + in (# goScale (j +# 1#) s3, () #) +{-# INLINE rawCholColumnSIMD #-} + +-- | Like 'rawCholColumnSIMD' but the dot-product starts from column @fromCol@ +-- instead of column 0. Used by panel Cholesky: after applying the GEMM update +-- from previous panels, the within-panel factorisation only needs contributions +-- from columns @fromCol .. j-1@. +rawCholColumnSIMDFrom :: MutableByteArray s -> Int -> Int -> Int -> Int -> ST s () +rawCholColumnSIMDFrom (MutableByteArray mba) (I# off) (I# n) (I# j) (I# fromCol) = ST $ \s0 -> + let jRowOff = off +# j *# n + dotLen = j -# fromCol + dotLen4 = dotLen -# (dotLen `remInt#` 4#) + k4End = fromCol +# dotLen4 + + goI i s + | isTrue# (i >=# n) = s + | otherwise = + let iRowOff = off +# i *# n + in case mutRowDot iRowOff jRowOff fromCol k4End j s of + (# s', dot #) -> + case readDoubleArray# mba (iRowOff +# j) s' of + (# s'', gij #) -> + case writeDoubleArray# mba (iRowOff +# j) (gij -## dot) s'' of + s''' -> goI (i +# 1#) s''' + + mutRowDot r1 r2 k kSimdEnd kEnd s + | isTrue# (k <# kSimdEnd) = + case goSimd r1 r2 k kSimdEnd (broadcastDoubleX4# 0.0##) s of + (# s', acc4 #) -> + let !(# a, b, c, d #) = unpackDoubleX4# acc4 + simdSum = a +## b +## c +## d + in mutRowDotScalar r1 r2 kSimdEnd kEnd simdSum s' + | otherwise = mutRowDotScalar r1 r2 k kEnd 0.0## s + + goSimd r1 r2 k kEnd acc s + | isTrue# (k >=# kEnd) = (# s, acc #) + | otherwise = + case readDoubleArrayAsDoubleX4# mba (r1 +# k) s of + (# s1, v1 #) -> + case readDoubleArrayAsDoubleX4# mba (r2 +# k) s1 of + (# s2, v2 #) -> goSimd r1 r2 (k +# 4#) kEnd (fmaddDoubleX4# v1 v2 acc) s2 + + mutRowDotScalar r1 r2 k kEnd acc s + | isTrue# (k >=# kEnd) = (# s, acc #) + | otherwise = + case readDoubleArray# mba (r1 +# k) s of + (# s1, v1 #) -> + case readDoubleArray# mba (r2 +# k) s1 of + (# s2, v2 #) -> mutRowDotScalar r1 r2 (k +# 1#) kEnd (acc +## v1 *## v2) s2 + + in case goI j s0 of + s1 -> + case readDoubleArray# mba (jRowOff +# j) s1 of + (# s2, gjj #) -> + let sjj = sqrtDouble# gjj + in case writeDoubleArray# mba (jRowOff +# j) sjj s2 of + s3 -> + let goScale i s + | isTrue# (i >=# n) = s + | otherwise = + case readDoubleArray# mba (off +# i *# n +# j) s of + (# s4, gij #) -> + case writeDoubleArray# mba (off +# i *# n +# j) (gij /## sjj) s4 of + s5 -> goScale (i +# 1#) s5 + in (# goScale (j +# 1#) s3, () #) +{-# INLINE rawCholColumnSIMDFrom #-} + +-- -------------------------------------------------------------------------- +-- SIMD forward/back substitution kernels (dot-product formulation) +-- -------------------------------------------------------------------------- + +-- | SIMD forward substitution (unit lower triangular, dot-product formulation). +-- Solves Ly = b where L has unit diagonal; b is already in mba_x. +-- For each row i: x[i] -= dot(L[i, 0..i-1], x[0..i-1]). +-- L row slices are contiguous in row-major storage → SIMD-friendly. +rawForwardSubUnitPackedSIMD :: ByteArray -> Int -> Int + -> MutableByteArray s -> Int -> ST s () +rawForwardSubUnitPackedSIMD (ByteArray ba_lu) (I# off_lu) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# n) = s + | otherwise = + let rowOff = off_lu +# i *# n + dotLen = i + d8End = dotLen -# (dotLen `remInt#` 8#) + d4End = dotLen -# (dotLen `remInt#` 4#) + in case goSimd8 rowOff 0# d8End (broadcastDoubleX4# 0.0##) (broadcastDoubleX4# 0.0##) s of + (# s1, acc0, acc1 #) -> + case goSimd4 rowOff d8End d4End acc0 s1 of + (# s2, acc0' #) -> + let !combined = plusDoubleX4# acc0' acc1 + !(# a, b, c, d #) = unpackDoubleX4# combined + simdSum = a +## b +## c +## d + in case goScalar rowOff d4End dotLen simdSum s2 of + (# s3, dotVal #) -> + case readDoubleArray# mba_x (off_x +# i) s3 of + (# s4, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) (xi -## dotVal) s4 of + s5 -> goI (i +# 1#) s5 + + goSimd8 rowOff k k8End acc0 acc1 s + | isTrue# (k >=# k8End) = (# s, acc0, acc1 #) + | otherwise = + let lv0 = indexDoubleArrayAsDoubleX4# ba_lu (rowOff +# k) + lv1 = indexDoubleArrayAsDoubleX4# ba_lu (rowOff +# k +# 4#) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# k) s of + (# s', xv0 #) -> + case readDoubleArrayAsDoubleX4# mba_x (off_x +# k +# 4#) s' of + (# s'', xv1 #) -> goSimd8 rowOff (k +# 8#) k8End + (fmaddDoubleX4# lv0 xv0 acc0) (fmaddDoubleX4# lv1 xv1 acc1) s'' + + goSimd4 rowOff k k4End acc s + | isTrue# (k >=# k4End) = (# s, acc #) + | otherwise = + let lv = indexDoubleArrayAsDoubleX4# ba_lu (rowOff +# k) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# k) s of + (# s', xv #) -> goSimd4 rowOff (k +# 4#) k4End (fmaddDoubleX4# lv xv acc) s' + + goScalar rowOff k kEnd acc s + | isTrue# (k >=# kEnd) = (# s, acc #) + | otherwise = + let lk = indexDoubleArray# ba_lu (rowOff +# k) + in case readDoubleArray# mba_x (off_x +# k) s of + (# s', xk #) -> goScalar rowOff (k +# 1#) kEnd (acc +## lk *## xk) s' + + in (# goI 0# s0, () #) +{-# INLINE rawForwardSubUnitPackedSIMD #-} + +-- | SIMD back substitution (upper triangular, dot-product formulation). +-- Solves Ux = y where y is in mba_x; overwrites with x. +-- For each row i (n-1 down to 0): x[i] = (x[i] - dot(U[i, i+1..n-1], x[i+1..n-1])) / U[i,i]. +rawBackSubPackedSIMD :: ByteArray -> Int -> Int + -> MutableByteArray s -> Int -> ST s () +rawBackSubPackedSIMD (ByteArray ba_lu) (I# off_lu) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goI i s + | isTrue# (i <# 0#) = s + | otherwise = + let rowOff = off_lu +# i *# n + dotStart = i +# 1# + dotLen = n -# i -# 1# + d8End = dotStart +# (dotLen -# (dotLen `remInt#` 8#)) + d4End = dotStart +# (dotLen -# (dotLen `remInt#` 4#)) + in case goSimd8 rowOff dotStart d8End (broadcastDoubleX4# 0.0##) (broadcastDoubleX4# 0.0##) s of + (# s1, acc0, acc1 #) -> + case goSimd4 rowOff d8End d4End acc0 s1 of + (# s2, acc0' #) -> + let !combined = plusDoubleX4# acc0' acc1 + !(# a, b, c, d #) = unpackDoubleX4# combined + simdSum = a +## b +## c +## d + in case goScalar rowOff d4End n simdSum s2 of + (# s3, dotVal #) -> + let uii = indexDoubleArray# ba_lu (rowOff +# i) + in case readDoubleArray# mba_x (off_x +# i) s3 of + (# s4, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) ((xi -## dotVal) /## uii) s4 of + s5 -> goI (i -# 1#) s5 + + goSimd8 rowOff k k8End acc0 acc1 s + | isTrue# (k >=# k8End) = (# s, acc0, acc1 #) + | otherwise = + let uv0 = indexDoubleArrayAsDoubleX4# ba_lu (rowOff +# k) + uv1 = indexDoubleArrayAsDoubleX4# ba_lu (rowOff +# k +# 4#) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# k) s of + (# s', xv0 #) -> + case readDoubleArrayAsDoubleX4# mba_x (off_x +# k +# 4#) s' of + (# s'', xv1 #) -> goSimd8 rowOff (k +# 8#) k8End + (fmaddDoubleX4# uv0 xv0 acc0) (fmaddDoubleX4# uv1 xv1 acc1) s'' + + goSimd4 rowOff k k4End acc s + | isTrue# (k >=# k4End) = (# s, acc #) + | otherwise = + let uv = indexDoubleArrayAsDoubleX4# ba_lu (rowOff +# k) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# k) s of + (# s', xv #) -> goSimd4 rowOff (k +# 4#) k4End (fmaddDoubleX4# uv xv acc) s' + + goScalar rowOff k kEnd acc s + | isTrue# (k >=# kEnd) = (# s, acc #) + | otherwise = + let uk = indexDoubleArray# ba_lu (rowOff +# k) + in case readDoubleArray# mba_x (off_x +# k) s of + (# s', xk #) -> goScalar rowOff (k +# 1#) kEnd (acc +## uk *## xk) s' + + in (# goI (n -# 1#) s0, () #) +{-# INLINE rawBackSubPackedSIMD #-} + +-- | SIMD Cholesky forward substitution (non-unit diagonal, dot-product formulation). +-- Solves Gy = b; b is in mba_x, overwrites with y. +-- For each row i: x[i] = (x[i] - dot(G[i, 0..i-1], x[0..i-1])) / G[i,i]. +rawForwardSubCholPackedSIMD :: ByteArray -> Int -> Int + -> MutableByteArray s -> Int -> ST s () +rawForwardSubCholPackedSIMD (ByteArray ba_g) (I# off_g) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# n) = s + | otherwise = + let rowOff = off_g +# i *# n + dotLen = i + d8End = dotLen -# (dotLen `remInt#` 8#) + d4End = dotLen -# (dotLen `remInt#` 4#) + in case goSimd8 rowOff 0# d8End (broadcastDoubleX4# 0.0##) (broadcastDoubleX4# 0.0##) s of + (# s1, acc0, acc1 #) -> + case goSimd4 rowOff d8End d4End acc0 s1 of + (# s2, acc0' #) -> + let !combined = plusDoubleX4# acc0' acc1 + !(# a, b, c, d #) = unpackDoubleX4# combined + simdSum = a +## b +## c +## d + in case goScalar rowOff d4End dotLen simdSum s2 of + (# s3, dotVal #) -> + let gii = indexDoubleArray# ba_g (rowOff +# i) + in case readDoubleArray# mba_x (off_x +# i) s3 of + (# s4, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) ((xi -## dotVal) /## gii) s4 of + s5 -> goI (i +# 1#) s5 + + goSimd8 rowOff k k8End acc0 acc1 s + | isTrue# (k >=# k8End) = (# s, acc0, acc1 #) + | otherwise = + let gv0 = indexDoubleArrayAsDoubleX4# ba_g (rowOff +# k) + gv1 = indexDoubleArrayAsDoubleX4# ba_g (rowOff +# k +# 4#) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# k) s of + (# s', xv0 #) -> + case readDoubleArrayAsDoubleX4# mba_x (off_x +# k +# 4#) s' of + (# s'', xv1 #) -> goSimd8 rowOff (k +# 8#) k8End + (fmaddDoubleX4# gv0 xv0 acc0) (fmaddDoubleX4# gv1 xv1 acc1) s'' + + goSimd4 rowOff k k4End acc s + | isTrue# (k >=# k4End) = (# s, acc #) + | otherwise = + let gv = indexDoubleArrayAsDoubleX4# ba_g (rowOff +# k) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# k) s of + (# s', xv #) -> goSimd4 rowOff (k +# 4#) k4End (fmaddDoubleX4# gv xv acc) s' + + goScalar rowOff k kEnd acc s + | isTrue# (k >=# kEnd) = (# s, acc #) + | otherwise = + let gk = indexDoubleArray# ba_g (rowOff +# k) + in case readDoubleArray# mba_x (off_x +# k) s of + (# s', xk #) -> goScalar rowOff (k +# 1#) kEnd (acc +## gk *## xk) s' + + in (# goI 0# s0, () #) +{-# INLINE rawForwardSubCholPackedSIMD #-} + +-- | SIMD Cholesky G^T back substitution (SAXPY formulation with broadcast). +-- Solves G^T x = y; y is in mba_x, overwrites with x. +-- For each j (n-1 down to 0): x[j] /= G[j,j], then for i=0..j-1: +-- x[i] -= G[j,i] * x[j] (SAXPY with broadcast x[j]). +-- G[j, 0..j-1] is contiguous in row-major → SIMD-friendly. +rawBackSubCholTPackedSIMD :: ByteArray -> Int -> Int + -> MutableByteArray s -> Int -> ST s () +rawBackSubCholTPackedSIMD (ByteArray ba_g) (I# off_g) (I# n) + (MutableByteArray mba_x) (I# off_x) = ST $ \s0 -> + let goJ j s + | isTrue# (j <# 0#) = s + | otherwise = + let gjj = indexDoubleArray# ba_g (off_g +# j *# n +# j) + in case readDoubleArray# mba_x (off_x +# j) s of + (# s', xj_ #) -> + let xj = xj_ /## gjj + in case writeDoubleArray# mba_x (off_x +# j) xj s' of + s'' -> + let jRowOff = off_g +# j *# n + negXj4 = broadcastDoubleX4# (negateDouble# xj) + updateLen = j + u8End = updateLen -# (updateLen `remInt#` 8#) + u4End = updateLen -# (updateLen `remInt#` 4#) + + goSimd8 i s_ + | isTrue# (i >=# u8End) = s_ + | otherwise = + let gv0 = indexDoubleArrayAsDoubleX4# ba_g (jRowOff +# i) + gv1 = indexDoubleArrayAsDoubleX4# ba_g (jRowOff +# i +# 4#) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# i) s_ of + (# s1, xv0 #) -> + case readDoubleArrayAsDoubleX4# mba_x (off_x +# i +# 4#) s1 of + (# s2, xv1 #) -> + case writeDoubleArrayAsDoubleX4# mba_x (off_x +# i) (fmaddDoubleX4# negXj4 gv0 xv0) s2 of + s3 -> case writeDoubleArrayAsDoubleX4# mba_x (off_x +# i +# 4#) (fmaddDoubleX4# negXj4 gv1 xv1) s3 of + s4 -> goSimd8 (i +# 8#) s4 + + goSimd4 i s_ + | isTrue# (i >=# u4End) = s_ + | otherwise = + let gv = indexDoubleArrayAsDoubleX4# ba_g (jRowOff +# i) + in case readDoubleArrayAsDoubleX4# mba_x (off_x +# i) s_ of + (# s1, xv #) -> + case writeDoubleArrayAsDoubleX4# mba_x (off_x +# i) (fmaddDoubleX4# negXj4 gv xv) s1 of + s2 -> goSimd4 (i +# 4#) s2 + + goScalar i s_ + | isTrue# (i >=# j) = s_ + | otherwise = + let gji = indexDoubleArray# ba_g (jRowOff +# i) + in case readDoubleArray# mba_x (off_x +# i) s_ of + (# s1, xi #) -> + case writeDoubleArray# mba_x (off_x +# i) (xi -## gji *## xj) s1 of + s2 -> goScalar (i +# 1#) s2 + + in goJ (j -# 1#) (goScalar u4End (goSimd4 u8End (goSimd8 0# s''))) + in (# goJ (n -# 1#) s0, () #) +{-# INLINE rawBackSubCholTPackedSIMD #-} + +-- -------------------------------------------------------------------------- +-- Utilities +-- -------------------------------------------------------------------------- + +minI :: Int# -> Int# -> Int# +minI a b = if isTrue# (a <=# b) then a else b +{-# INLINE minI #-} diff --git a/src/Numeric/LinearAlgebra/Massiv/Linear.hs b/src/Numeric/LinearAlgebra/Massiv/Linear.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Linear.hs @@ -0,0 +1,122 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Linear +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = Integration with the @linear@ library +-- +-- This module provides conversion functions between the +-- library's types +-- (@'Linear.V.V'@, @'Linear.V2.V2'@, @'Linear.V3.V3'@, @'Linear.V4.V4'@) +-- and our dimensioned @'Vector'@ \/ @'Matrix'@ types. +-- +-- == Why not typeclass instances? +-- +-- The @linear@ library's typeclasses ('Linear.Additive.Additive', +-- 'Linear.Metric.Metric', 'Linear.Trace.Trace') expect types of kind +-- @* -> *@ (i.e., functors over the element type). Our @Vector n r e@ and +-- @Matrix m n r e@ carry additional type parameters (@n@, @r@) before @e@, +-- making direct functor-based instances impractical without additional +-- newtype wrappers. +-- +-- Instead, equivalent operations are provided as standalone functions: +-- +-- * "Numeric.LinearAlgebra.Massiv.BLAS.Level1" — 'dot', 'axpy', 'scal', 'nrm2' +-- * "Numeric.LinearAlgebra.Massiv.BLAS.Level3" — 'mAdd', 'mSub', 'mScale', 'matMul', 'transpose' +-- * "Numeric.LinearAlgebra.Massiv.Norms" — 'normFrob', 'norm1', 'normInf' +-- +-- == Conversion semantics +-- +-- 'fromLinearV' converts to the @'Data.Massiv.Array.B'@ (boxed) representation +-- because @linear@'s @V@ stores elements in a boxed @Data.Vector.Vector@. +-- For small fixed-size types (@V2@, @V3@, @V4@), 'fromV2' etc. produce +-- vectors in any representation @r@ satisfying @Manifest r e@. +module Numeric.LinearAlgebra.Massiv.Linear + ( -- * Conversions with @linear@'s @V n a@ + fromLinearV + , toLinearV + -- * Small fixed-size vector conversions + , fromV2 + , fromV3 + , fromV4 + -- * List-based matrix I\/O + , toListMatrix + , fromListMatrix + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..)) +import qualified Data.Vector as BV +import GHC.TypeNats (KnownNat, natVal) +import Data.Proxy (Proxy(..)) + +import qualified Linear.V as L +import Linear.V2 (V2(..)) +import Linear.V3 (V3(..)) +import Linear.V4 (V4(..)) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal + +-- | Convert a @linear@ @'Linear.V.V' n a@ to our @'Vector' n 'Data.Massiv.Array.B' a@. +-- +-- The result uses the boxed (@B@) representation since @linear@'s @V@ is +-- backed by a boxed @Data.Vector.Vector@. +fromLinearV :: forall n a. KnownNat n => L.V n a -> Vector n M.B a +fromLinearV lv = + let bv = L.toVector lv + nn = fromIntegral (natVal (Proxy @n)) + arr = M.compute @M.B $ M.makeArray @M.D M.Seq (M.Sz1 nn) (bv BV.!) + in MkVector arr + +-- | Convert our @'Vector' n 'Data.Massiv.Array.B' a@ to a @linear@ @'Linear.V.V' n a@. +-- +-- This is the inverse of 'fromLinearV'. The dimension is checked by @linear@'s +-- 'Linear.V.fromVector' (which returns 'Maybe'); the 'error' case is unreachable +-- given correct type-level dimensions. +toLinearV :: forall n a. KnownNat n => Vector n M.B a -> L.V n a +toLinearV (MkVector arr) = + let nn = fromIntegral (natVal (Proxy @n)) + bv = BV.generate nn (\i -> M.index' arr i) + in case L.fromVector bv of + Just v -> v + Nothing -> error "toLinearV: impossible dimension mismatch" + +-- | Convert a @linear@ @'Linear.V2.V2'@ to a 2-element 'Vector'. +fromV2 :: M.Manifest r e => V2 e -> Vector 2 r e +fromV2 (V2 x y) = makeVector @2 $ \i -> case i of { 0 -> x; _ -> y } + +-- | Convert a @linear@ @'Linear.V3.V3'@ to a 3-element 'Vector'. +fromV3 :: M.Manifest r e => V3 e -> Vector 3 r e +fromV3 (V3 x y z) = makeVector @3 $ \i -> case i of { 0 -> x; 1 -> y; _ -> z } + +-- | Convert a @linear@ @'Linear.V4.V4'@ to a 4-element 'Vector'. +fromV4 :: M.Manifest r e => V4 e -> Vector 4 r e +fromV4 (V4 x y z w) = makeVector @4 $ \i -> case i of { 0 -> x; 1 -> y; 2 -> z; _ -> w } + +-- | Convert a matrix to a list of lists (row-major order). +-- +-- @ +-- toListMatrix mat == [[mat '!' (i,j) | j <- [0..n-1]] | i <- [0..m-1]] +-- @ +toListMatrix :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e) + => Matrix m n r e -> [[e]] +toListMatrix mat = + let r = dimVal @m + c = dimVal @n + in [[mat ! (i, j) | j <- [0..c-1]] | i <- [0..r-1]] + +-- | Create a matrix from a list of lists (row-major order). +-- +-- Returns 'Nothing' if the list dimensions do not match the type-level +-- dimensions \(m\) and \(n\). +fromListMatrix :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e) + => [[e]] -> Maybe (Matrix m n r e) +fromListMatrix rows_ + | length rows_ /= dimVal @m = Nothing + | any (\row -> length row /= dimVal @n) rows_ = Nothing + | otherwise = Just $ makeMatrix @m @n @r $ \i j -> (rows_ !! i) !! j diff --git a/src/Numeric/LinearAlgebra/Massiv/Norms.hs b/src/Numeric/LinearAlgebra/Massiv/Norms.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Norms.hs @@ -0,0 +1,163 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Norms +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = Matrix and Vector Norms +-- +-- Norms measure the "size" of vectors and matrices and are fundamental to +-- error analysis, convergence criteria, and conditioning estimates in +-- numerical linear algebra. +-- +-- This module implements the norms described in: +-- +-- * Golub, G. H. & Van Loan, C. F. (2013). /Matrix Computations/, 4th ed. +-- __Chapter 2: Matrix Analysis__, Sections 2.3–2.7, pp. 71–95. +-- +-- == Vector norms (GVL4 Section 2.3, p. 71) +-- +-- For a vector \(x \in \mathbb{R}^n\): +-- +-- * 1-norm: \(\|x\|_1 = \sum_{i=1}^{n} |x_i|\) — see 'vnorm1' +-- * 2-norm (Euclidean): \(\|x\|_2 = \sqrt{\sum_{i=1}^{n} x_i^2}\) — see 'vnorm2' +-- * \(\infty\)-norm: \(\|x\|_\infty = \max_i |x_i|\) — see 'vnormInf' +-- +-- == Matrix norms (GVL4 Section 2.3, pp. 71–78) +-- +-- For a matrix \(A \in \mathbb{R}^{m \times n}\): +-- +-- * Frobenius norm: \(\|A\|_F = \sqrt{\sum_{i,j} a_{ij}^2} = \sqrt{\text{trace}(A^T A)}\) +-- — see 'normFrob' +-- * 1-norm (max column sum): \(\|A\|_1 = \max_j \sum_i |a_{ij}|\) — see 'norm1' +-- * \(\infty\)-norm (max row sum): \(\|A\|_\infty = \max_i \sum_j |a_{ij}|\) — see 'normInf' +-- +-- These satisfy the norm axioms: non-negativity, homogeneity, and the +-- triangle inequality \(\|A + B\| \leq \|A\| + \|B\|\). +-- +-- == Condition numbers (GVL4 Section 2.7, pp. 87–95) +-- +-- The condition number \(\kappa(A) = \|A\| \cdot \|A^{-1}\|\) measures +-- how sensitive the solution of \(Ax = b\) is to perturbations in \(A\) +-- and \(b\). See 'condFrob' for a placeholder using the Frobenius norm. +module Numeric.LinearAlgebra.Massiv.Norms + ( -- * Vector norms (GVL4 Section 2.3, p. 71) + vnorm1 + , vnorm2 + , vnormInf + -- * Matrix norms (GVL4 Section 2.3, pp. 71–78) + , normFrob + , norm1 + , normInf + -- * Condition number estimate (GVL4 Section 2.7, p. 87) + , condFrob + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal + +-- | Vector 1-norm (GVL4 p. 71, eq. 2.3.1). +-- +-- \[ +-- \|x\|_1 = \sum_{i=1}^{n} |x_i| +-- \] +-- +-- Complexity: \(O(n)\). +vnorm1 :: (KnownNat n, M.Manifest r e, Num e, Ord e) + => Vector n r e -> e +vnorm1 (MkVector arr) = M.foldlS (\acc x -> acc + abs x) 0 arr + +-- | Vector 2-norm, the Euclidean norm (GVL4 p. 71, eq. 2.3.2). +-- +-- \[ +-- \|x\|_2 = \sqrt{\sum_{i=1}^{n} x_i^2} = \sqrt{x^T x} +-- \] +-- +-- Complexity: \(O(n)\). +vnorm2 :: (KnownNat n, M.Manifest r e, Floating e) + => Vector n r e -> e +vnorm2 (MkVector arr) = sqrt $ M.foldlS (\acc x -> acc + x * x) 0 arr + +-- | Vector \(\infty\)-norm (GVL4 p. 71, eq. 2.3.3). +-- +-- \[ +-- \|x\|_\infty = \max_{1 \leq i \leq n} |x_i| +-- \] +-- +-- Complexity: \(O(n)\). +vnormInf :: (KnownNat n, M.Manifest r e, Num e, Ord e) + => Vector n r e -> e +vnormInf (MkVector arr) = M.foldlS (\acc x -> max acc (abs x)) 0 arr + +-- | Frobenius norm (GVL4 p. 72, eq. 2.3.7). +-- +-- \[ +-- \|A\|_F = \sqrt{\sum_{i=1}^{m} \sum_{j=1}^{n} a_{ij}^2} +-- = \sqrt{\text{trace}(A^T A)} +-- \] +-- +-- The Frobenius norm is /not/ an operator norm (it is not subordinate +-- to any vector norm), but it is submultiplicative: +-- \(\|AB\|_F \leq \|A\|_F \|B\|_F\). +-- +-- Complexity: \(O(mn)\). +normFrob :: (KnownNat m, KnownNat n, M.Manifest r e, Floating e) + => Matrix m n r e -> e +normFrob (MkMatrix arr) = sqrt $ M.foldlS (\acc x -> acc + x * x) 0 arr + +-- | Matrix 1-norm — maximum absolute column sum (GVL4 p. 72, eq. 2.3.10). +-- +-- \[ +-- \|A\|_1 = \max_{1 \leq j \leq n} \sum_{i=1}^{m} |a_{ij}| +-- \] +-- +-- This is the operator norm subordinate to the vector 1-norm: +-- \(\|A\|_1 = \max_{\|x\|_1 = 1} \|Ax\|_1\). +-- +-- Complexity: \(O(mn)\). +norm1 :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e, Ord e) + => Matrix m n r e -> e +norm1 mat = + let c = dimVal @n + r = dimVal @m + colSum j = foldl (\acc i -> acc + abs (mat ! (i, j))) 0 [0..r-1] + in maximum $ map colSum [0..c-1] + +-- | Matrix \(\infty\)-norm — maximum absolute row sum (GVL4 p. 72, eq. 2.3.11). +-- +-- \[ +-- \|A\|_\infty = \max_{1 \leq i \leq m} \sum_{j=1}^{n} |a_{ij}| +-- \] +-- +-- This is the operator norm subordinate to the vector \(\infty\)-norm. +-- Note that \(\|A\|_\infty = \|A^T\|_1\). +-- +-- Complexity: \(O(mn)\). +normInf :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e, Ord e) + => Matrix m n r e -> e +normInf mat = + let c = dimVal @n + r = dimVal @m + rowSum i = foldl (\acc j -> acc + abs (mat ! (i, j))) 0 [0..c-1] + in maximum $ map rowSum [0..r-1] + +-- | Estimate condition number using the Frobenius norm (GVL4 Section 2.7, p. 87). +-- +-- The condition number is defined as +-- \(\kappa_F(A) = \|A\|_F \cdot \|A^{-1}\|_F\). +-- +-- __Note__: This function currently returns only \(\|A\|_F\) as a placeholder. +-- Computing \(\|A^{-1}\|_F\) requires solving a linear system (e.g., via LU), +-- introducing a circular dependency. Users should compute the full condition +-- number by combining 'normFrob' with an explicit inverse or using SVD-based +-- estimates (\(\kappa_2 = \sigma_{\max} / \sigma_{\min}\)). +condFrob :: (KnownNat n, M.Manifest r e, Floating e) + => Matrix n n r e -> e +condFrob = normFrob diff --git a/src/Numeric/LinearAlgebra/Massiv/Orthogonal/Givens.hs b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/Givens.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/Givens.hs @@ -0,0 +1,156 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Orthogonal.Givens +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Givens rotations for selective zeroing of matrix entries. +-- +-- This module implements Givens (plane) rotations following Golub & Van +-- Loan, /Matrix Computations/, 4th edition (GVL4), Section 5.1.8, +-- pp. 240--243. +-- +-- A Givens rotation is an orthogonal matrix that operates in a +-- two-dimensional subspace. Given scalars \( a \) and \( b \), the +-- rotation matrix +-- +-- \( G^T = \begin{bmatrix} c & -s \\ s & c \end{bmatrix} \) +-- +-- is constructed so that +-- +-- \( G^T \begin{bmatrix} a \\ b \end{bmatrix} = \begin{bmatrix} r \\ 0 \end{bmatrix} \) +-- +-- where \( r = \sqrt{a^2 + b^2} \). Our convention follows GVL4 +-- Algorithm 5.1.3 (p. 240): \( c = a / r \), \( s = -b / r \). +-- +-- Givens rotations are especially useful when only a small number of +-- sub-diagonal entries need to be zeroed (e.g., in Hessenberg or banded +-- matrices), whereas Householder reflections are preferred for zeroing +-- entire sub-columns at once. Givens-based QR factorisation is the +-- method of choice for tridiagonal and Hessenberg eigenvalue problems +-- (GVL4 Section 5.2.8, p. 255). +-- +-- __Complexity.__ +-- +-- * Computing a Givens rotation ('givensRotation'): \( O(1) \) flops +-- (one square root and a small number of divisions). +-- * Applying a Givens rotation to a row or column pair of an +-- \( m \times n \) matrix ('applyGivensLeft', 'applyGivensRight'): +-- \( O(n) \) or \( O(m) \) flops respectively (one pass over the +-- affected row or column pair). +module Numeric.LinearAlgebra.Massiv.Orthogonal.Givens + ( -- * Givens rotation + givensRotation + -- * Apply Givens rotation + , applyGivensLeft + , applyGivensRight + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal + +-- | Compute a Givens rotation (GVL4 Algorithm 5.1.3, p. 240). +-- +-- Given scalars \( a \) and \( b \), compute cosine \( c \) and sine +-- \( s \) such that +-- +-- \( \begin{bmatrix} c & s \\ -s & c \end{bmatrix}^T \begin{bmatrix} a \\ b \end{bmatrix} = \begin{bmatrix} r \\ 0 \end{bmatrix} \) +-- +-- where \( r = \sqrt{a^2 + b^2} \). +-- +-- The implementation avoids overflow and unnecessary computation by +-- distinguishing three cases: +-- +-- * If \( b = 0 \), the rotation is the identity: \( c = 1, s = 0 \). +-- * If \( |b| > |a| \), the tangent \( \tau = -a/b \) is computed first, +-- then \( s = 1 / \sqrt{1 + \tau^2} \) and \( c = s \tau \). +-- * Otherwise, \( \tau = -b/a \), \( c = 1 / \sqrt{1 + \tau^2} \), and +-- \( s = c \tau \). +-- +-- This avoids computing the potentially large quantity +-- \( r = \sqrt{a^2 + b^2} \) directly, which could overflow. +-- +-- __Complexity:__ \( O(1) \) flops (one square root, a few multiplications +-- and divisions). +-- +-- Returns @(c, s)@. +givensRotation :: (Floating e, Ord e) => e -> e -> (e, e) +givensRotation a b + | b == 0 = (1, 0) + | abs b > abs a = + let tau = -a / b + s = 1 / sqrt (1 + tau * tau) + c = s * tau + in (c, s) + | otherwise = + let tau = -b / a + c = 1 / sqrt (1 + tau * tau) + s = c * tau + in (c, s) + +-- | Apply a Givens rotation from the left to rows @i@ and @k@ of a matrix +-- (GVL4 Section 5.1.9, p. 241). +-- +-- Performs the update +-- +-- \( A([i,k], :) \leftarrow G^T \, A([i,k], :) \) +-- +-- where \( G^T = \begin{bmatrix} c & -s \\ s & c \end{bmatrix} \). +-- Only rows @i@ and @k@ are modified; all other rows are untouched. +-- This is the standard operation used to zero out the \( (k, j) \) +-- entry of a matrix during Givens-based QR factorisation +-- (GVL4 Algorithm 5.2.3, p. 252). +-- +-- __Complexity:__ \( O(n) \) flops, where \( n \) is the number of +-- columns. +applyGivensLeft :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => e -- ^ c + -> e -- ^ s + -> Int -- ^ row i + -> Int -- ^ row k + -> Matrix m n r e -> Matrix m n r e +applyGivensLeft c s ri rk a = + makeMatrix @m @n @r $ \i j -> + if i == ri then + c * (a ! (ri, j)) - s * (a ! (rk, j)) + else if i == rk then + s * (a ! (ri, j)) + c * (a ! (rk, j)) + else + a ! (i, j) + +-- | Apply a Givens rotation from the right to columns @i@ and @k@ of a +-- matrix (GVL4 Section 5.1.9, p. 242). +-- +-- Performs the update +-- +-- \( A(:, [i,k]) \leftarrow A(:, [i,k]) \, G \) +-- +-- where \( G = \begin{bmatrix} c & s \\ -s & c \end{bmatrix} \). +-- Only columns @i@ and @k@ are modified; all other columns are +-- untouched. Right-multiplication by a Givens rotation is typically +-- used to accumulate the orthogonal factor \( Q \) during QR +-- factorisation (GVL4 Section 5.1.9, p. 242). +-- +-- __Complexity:__ \( O(m) \) flops, where \( m \) is the number of +-- rows. +applyGivensRight :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => e -- ^ c + -> e -- ^ s + -> Int -- ^ column i + -> Int -- ^ column k + -> Matrix m n r e -> Matrix m n r e +applyGivensRight c s ci ck a = + makeMatrix @m @n @r $ \i j -> + if j == ci then + c * (a ! (i, ci)) - s * (a ! (i, ck)) + else if j == ck then + s * (a ! (i, ci)) + c * (a ! (i, ck)) + else + a ! (i, j) diff --git a/src/Numeric/LinearAlgebra/Massiv/Orthogonal/Householder.hs b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/Householder.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/Householder.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Orthogonal.Householder +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Householder reflections for orthogonal triangularisation. +-- +-- This module implements the Householder reflection (also known as a +-- Householder transformation), following Golub & Van Loan, /Matrix +-- Computations/, 4th edition (GVL4), Section 5.1, pp. 236--243. +-- +-- As GVL4 states (p. 236): \"The Householder reflection is the most +-- important tool in matrix computations.\" A Householder reflector is a +-- matrix of the form +-- +-- \( P = I - \beta v v^T \) +-- +-- where \( v \) is the /Householder vector/ and \( \beta = 2 / (v^T v) \). +-- The key property of \( P \) is that it is both symmetric and orthogonal: +-- +-- \( P = P^T = P^{-1} \) +-- +-- Given an input vector \( x \), the Householder vector \( v \) and scalar +-- \( \beta \) are chosen so that +-- +-- \( P x = (I - \beta v v^T) x = \| x \|_2 \, e_1 \) +-- +-- where \( e_1 \) is the first standard basis vector. This is the +-- fundamental operation behind Householder QR factorisation (GVL4 +-- Algorithm 5.2.1) and many other matrix decompositions. +-- +-- __Complexity.__ +-- +-- * Computing the Householder vector ('householderVector'): \( O(n) \) flops. +-- * Applying a Householder reflection to an \( m \times n \) matrix +-- ('applyHouseholderLeft', 'applyHouseholderRight'): \( O(mn) \) flops. +-- * Forming the explicit reflector matrix ('householderMatrix'): \( O(n^2) \) +-- flops, but this should be avoided in favour of implicit application +-- whenever possible. +module Numeric.LinearAlgebra.Massiv.Orthogonal.Householder + ( -- * Householder vector + householderVector + -- * Apply Householder reflection + , applyHouseholderLeft + , applyHouseholderRight + -- * Construct explicit reflector + , householderMatrix + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal + +-- | Compute a Householder vector (GVL4 Algorithm 5.1.1, p. 236). +-- +-- Given a vector \( x \in \mathbb{R}^n \), compute the Householder vector +-- \( v \) and scalar \( \beta \) such that +-- +-- \( (I - \beta \, v \, v^T) \, x = \| x \|_2 \, e_1 \) +-- +-- where \( e_1 \) is the first standard basis vector. By convention the +-- first component of \( v \) is normalised to \( v_1 = 1 \), which allows +-- it to be stored implicitly in the sub-diagonal part of a matrix during +-- QR factorisation. +-- +-- The implementation follows GVL4 Algorithm 5.1.1 exactly, including the +-- careful treatment of the sign of \( x_1 \) to avoid catastrophic +-- cancellation. When \( x \) is already a non-negative multiple of +-- \( e_1 \), the function returns \( \beta = 0 \) (i.e., the identity +-- transformation). +-- +-- __Complexity:__ \( O(n) \) flops. +-- +-- Returns @(v, beta)@. +householderVector :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Vector n r e -> (Vector n r e, e) +householderVector x = + let nn = dimVal @n + x0 = x !. 0 + -- σ = x(2:n)ᵀ · x(2:n) + sigma = foldl' (\acc i -> acc + (x !. i) * (x !. i)) 0 [1..nn-1] + in if sigma == 0 && x0 >= 0 + then -- x is already a positive multiple of e1 + ( makeVector @n @r $ \i -> if i == 0 then 1 else 0 + , 0 + ) + else if sigma == 0 + then -- x = -α·e₁ + ( makeVector @n @r $ \i -> if i == 0 then 1 else 0 + , 2 + ) + else + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + v = makeVector @n @r $ \i -> + if i == 0 then 1 else (x !. i) / v0 + in (v, beta) + +-- | Apply a Householder reflection from the left (GVL4 Section 5.1, p. 236). +-- +-- Given a Householder vector \( v \in \mathbb{R}^m \), scalar \( \beta \), +-- and matrix \( A \in \mathbb{R}^{m \times n} \), compute +-- +-- \( A \leftarrow (I - \beta \, v \, v^T) \, A = A - \beta \, v \, (A^T v)^T \) +-- +-- The computation is performed without forming \( P \) explicitly. +-- Instead, the intermediate vector \( w = \beta \, A^T v \) is computed +-- first, and then the rank-1 update \( A \leftarrow A - v \, w^T \) is +-- applied. This is the standard technique described in GVL4 +-- Section 5.1.4 (p. 238). +-- +-- __Complexity:__ \( O(mn) \) flops. +applyHouseholderLeft :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => Vector m r e -> e -> Matrix m n r e -> Matrix m n r e +applyHouseholderLeft v beta a = + let mm = dimVal @m + c = dimVal @n + in makeMatrix @m @n @r $ \i j -> + let -- w = βAᵀv, w(j) = β · Σᵢ v(i)·A(i,j) + wj = beta * foldl' (\acc k -> acc + (v !. k) * (a ! (k, j))) 0 [0..mm-1] + in (a ! (i, j)) - (v !. i) * wj + +-- | Apply a Householder reflection from the right (GVL4 Section 5.1, p. 236). +-- +-- Given a matrix \( A \in \mathbb{R}^{m \times n} \), Householder vector +-- \( v \in \mathbb{R}^n \), and scalar \( \beta \), compute +-- +-- \( A \leftarrow A \, (I - \beta \, v \, v^T) = A - \beta \, (A \, v) \, v^T \) +-- +-- As with 'applyHouseholderLeft', the reflector is never formed +-- explicitly. The intermediate vector \( w = \beta \, A \, v \) is +-- computed first, followed by the rank-1 update +-- \( A \leftarrow A - w \, v^T \). See GVL4 Section 5.1.4 (p. 238). +-- +-- __Complexity:__ \( O(mn) \) flops. +applyHouseholderRight :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Num e) + => Matrix m n r e -> Vector n r e -> e -> Matrix m n r e +applyHouseholderRight a v beta = + let c = dimVal @n + in makeMatrix @m @n @r $ \i j -> + let -- w = β·A·v, w(i) = β · Σⱼ A(i,j)·v(j) + wi = beta * foldl' (\acc k -> acc + (a ! (i, k)) * (v !. k)) 0 [0..c-1] + in (a ! (i, j)) - wi * (v !. j) + +-- | Construct the explicit Householder reflector matrix (GVL4 Section 5.1, p. 236). +-- +-- Given a Householder vector \( v \in \mathbb{R}^n \) and scalar +-- \( \beta \), form the \( n \times n \) matrix +-- +-- \( H = I - \beta \, v \, v^T \) +-- +-- The resulting matrix is both symmetric and orthogonal: +-- \( H = H^T = H^{-1} \). +-- +-- __Note:__ In most numerical algorithms it is preferable to apply the +-- Householder transformation implicitly via 'applyHouseholderLeft' or +-- 'applyHouseholderRight' rather than forming \( H \) explicitly. +-- Forming the explicit matrix costs \( O(n^2) \) flops and storage, and +-- subsequent multiplication with it costs \( O(n^3) \) rather than the +-- \( O(mn) \) achievable by implicit application. +-- +-- __Complexity:__ \( O(n^2) \) flops. +householderMatrix :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Vector n r e -> e -> Matrix n n r e +householderMatrix v beta = + makeMatrix @n @n @r $ \i j -> + let ident = if i == j then 1 else 0 + in ident - beta * (v !. i) * (v !. j) diff --git a/src/Numeric/LinearAlgebra/Massiv/Orthogonal/LeastSquares.hs b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/LeastSquares.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/LeastSquares.hs @@ -0,0 +1,133 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Orthogonal.LeastSquares +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Linear least squares solvers. +-- +-- This module provides two methods for solving the linear least squares +-- problem +-- +-- \( \min_x \| A x - b \|_2 \) +-- +-- following Golub & Van Loan, /Matrix Computations/, 4th edition (GVL4), +-- Section 5.3, pp. 260--270. +-- +-- __Theorem 5.3.1 (Least squares existence, GVL4 p. 260).__ Let +-- \( A \in \mathbb{R}^{m \times n} \) with \( m \ge n \) and +-- \( \operatorname{rank}(A) = n \). Then the least squares problem +-- \( \min_x \| A x - b \|_2 \) has a unique solution \( x^* \) given +-- by the /normal equations/ +-- +-- \( A^T A \, x = A^T b \) +-- +-- Two solution methods are provided: +-- +-- * __QR-based__ ('leastSquaresQR') -- GVL4 Algorithm 5.3.2 (p. 262). +-- Factor \( A = Q R \) via Householder QR, then solve +-- \( R_1 x = (Q^T b)_{1:n} \) by back-substitution, where \( R_1 \) +-- denotes the leading \( n \times n \) upper triangular block of +-- \( R \). This is the recommended method: it is numerically stable +-- and does not square the condition number. +-- +-- * __Normal equations__ ('leastSquaresNormal') -- GVL4 Section 5.3.2 +-- (p. 261). Form \( A^T A \) and \( A^T b \) explicitly, then solve +-- via Cholesky factorisation. This is faster but squares the +-- condition number: \( \kappa_2(A^T A) = \kappa_2(A)^2 \), so it +-- should only be used when \( A \) is well-conditioned. +-- +-- __Complexity.__ +-- +-- * QR-based: \( 2mn^2 \) flops (dominated by the QR factorisation). +-- * Normal equations: \( mn^2 + \tfrac{1}{3}n^3 \) flops +-- (\( mn^2 \) for forming \( A^T A \), \( \tfrac{1}{3}n^3 \) for +-- Cholesky). +module Numeric.LinearAlgebra.Massiv.Orthogonal.LeastSquares + ( -- * Least squares via QR + leastSquaresQR + -- * Normal equations + , leastSquaresNormal + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Orthogonal.QR (qr) +import Numeric.LinearAlgebra.Massiv.Solve.Triangular (backSub) +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, transpose) +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvec) +import Numeric.LinearAlgebra.Massiv.Solve.Cholesky (choleskySolve) + +-- | Solve the least squares problem \( \min_x \| A x - b \|_2 \) via QR +-- factorisation (GVL4 Algorithm 5.3.2, p. 262). +-- +-- Given \( A \in \mathbb{R}^{m \times n} \) with \( m \ge n \) and full +-- column rank, and a right-hand side \( b \in \mathbb{R}^m \), compute +-- the unique least squares solution \( x^* \in \mathbb{R}^n \). +-- +-- The algorithm proceeds in three steps: +-- +-- 1. Factor \( A = Q R \) using Householder QR ('qr'). +-- 2. Form the transformed right-hand side \( Q^T b \). +-- 3. Solve the \( n \times n \) upper triangular system +-- \( R_1 x = (Q^T b)_{1:n} \) by back-substitution, where \( R_1 \) +-- is the leading \( n \times n \) block of \( R \). +-- +-- This method is numerically stable because orthogonal transformations +-- preserve the 2-norm and do not amplify rounding errors. The condition +-- number relevant to the solution is \( \kappa_2(A) \), not +-- \( \kappa_2(A)^2 \) as with the normal equations. +-- +-- __Complexity:__ \( 2mn^2 - \tfrac{2}{3}n^3 \) flops for the QR +-- factorisation plus \( 2mn \) flops for forming \( Q^T b \) and +-- \( n^2 \) flops for back-substitution, giving a total of +-- \( O(2mn^2) \) flops. +leastSquaresQR :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix m n r e -> Vector m r e -> Vector n r e +leastSquaresQR a b = + let mm = dimVal @m + nn = dimVal @n + (q, r) = qr a + qt = transpose q + -- qtb = Qᵀ·b (dimension m) + qtb = matvec qt b + -- Extract top n×n of R (which is upper triangular) + r1 = makeMatrix @n @n @r $ \i j -> r ! (i, j) + -- Extract top n entries of Qᵀb + qtb1 = makeVector @n @r $ \i -> qtb !. i + in backSub r1 qtb1 + +-- | Solve the least squares problem \( \min_x \| A x - b \|_2 \) via the +-- normal equations (GVL4 Section 5.3.2, p. 261). +-- +-- Given \( A \in \mathbb{R}^{m \times n} \) with full column rank and +-- \( b \in \mathbb{R}^m \), form the \( n \times n \) symmetric positive +-- definite system +-- +-- \( A^T A \, x = A^T b \) +-- +-- and solve it using Cholesky factorisation (\( A^T A = L L^T \)). +-- +-- __Warning:__ The normal equations square the condition number of \( A \): +-- \( \kappa_2(A^T A) = \kappa_2(A)^2 \). For ill-conditioned problems +-- this leads to a significant loss of accuracy compared to QR-based +-- methods. Prefer 'leastSquaresQR' unless \( A \) is known to be +-- well-conditioned (GVL4 p. 261). +-- +-- __Complexity:__ \( mn^2 \) flops for forming \( A^T A \), plus +-- \( \tfrac{1}{3}n^3 \) flops for the Cholesky factorisation, giving a +-- total of \( O(mn^2 + \tfrac{1}{3}n^3) \) flops. +leastSquaresNormal :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix m n r e -> Vector m r e -> Vector n r e +leastSquaresNormal a b = + let at = transpose a + ata = matMul at a -- n×n, symmetric positive definite + atb = matvec at b -- n×1 + in choleskySolve ata atb diff --git a/src/Numeric/LinearAlgebra/Massiv/Orthogonal/QR.hs b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/QR.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Orthogonal/QR.hs @@ -0,0 +1,481 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Orthogonal.QR +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- QR factorisation via Householder reflections and Givens rotations. +-- +-- This module implements the QR factorisation of a matrix +-- \( A \in \mathbb{R}^{m \times n} \) with \( m \ge n \), following +-- Golub & Van Loan, /Matrix Computations/, 4th edition (GVL4), +-- Section 5.2, pp. 246--260. +-- +-- __Theorem 5.2.1 (QR existence, GVL4 p. 246).__ For any +-- \( A \in \mathbb{R}^{m \times n} \) with \( m \ge n \), there exists +-- an orthogonal matrix \( Q \in \mathbb{R}^{m \times m} \) and an upper +-- triangular matrix \( R \in \mathbb{R}^{m \times n} \) such that +-- +-- \( A = Q \, R \) +-- +-- If \( A \) has full column rank, the factorisation is unique up to +-- sign changes in the rows of \( R \) (equivalently, columns of \( Q \)). +-- +-- Two algorithms are provided: +-- +-- * __Householder QR__ ('qr', 'qrR') -- GVL4 Algorithm 5.2.1 (p. 249). +-- A sequence of Householder reflections \( H_1, H_2, \ldots, H_n \) +-- is applied to \( A \) from the left to produce +-- \( H_n \cdots H_2 \, H_1 \, A = R \), so that +-- \( Q = H_1 \, H_2 \cdots H_n \). +-- +-- * __Givens QR__ ('qrGivens') -- GVL4 Section 5.2.4 (p. 252). +-- A sequence of Givens rotations zeroes out sub-diagonal entries one +-- at a time. This variant is preferred for sparse or banded matrices, +-- particularly Hessenberg matrices, where the number of rotations is +-- proportional to the bandwidth rather than the matrix dimension. +-- +-- __Complexity.__ +-- +-- * Householder QR: \( 2mn^2 - \tfrac{2}{3}n^3 \) flops (GVL4 p. 249). +-- * Givens QR: \( 3mn^2 - n^3 \) flops for a dense matrix, but +-- significantly fewer for structured (e.g., Hessenberg) matrices. +-- +-- __Optimisation.__ The implementation uses in-place mutable arrays via +-- the 'ST' monad, storing Householder vectors implicitly in the +-- subdiagonal of the working matrix (the LAPACK convention). The +-- orthogonal factor \( Q \) is formed via backward accumulation, and +-- 'qrR' avoids forming \( Q \) entirely. +module Numeric.LinearAlgebra.Massiv.Orthogonal.QR + ( -- * QR factorisation (Householder) + qr + , qrP + , qrR + -- * QR factorisation (Givens) + , qrGivens + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..), unwrapByteArray, unwrapByteArrayOffset, + unwrapMutableByteArray, unwrapMutableByteArrayOffset) +import GHC.TypeNats (KnownNat) +import Control.Monad (when) +import Control.Monad.ST (ST) +import GHC.Exts +import GHC.ST (ST(..)) +import Data.Primitive.ByteArray (ByteArray(..), MutableByteArray(..), newByteArray, unsafeFreezeByteArray) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Orthogonal.Givens + (givensRotation) +import Numeric.LinearAlgebra.Massiv.Internal.Kernel + (rawMutSumSqColumn, rawMutSumProdColumns, rawMutHouseholderApply, rawMutQAccum, + rawGemmKernel, rawZeroDoubles, rawNegateDoubles) + +-- | Full QR factorisation via Householder reflections (GVL4 Algorithm 5.2.1, p. 249). +-- +-- Given \( A \in \mathbb{R}^{m \times n} \) with \( m \ge n \), compute +-- the factorisation +-- +-- \( A = Q \, R \) +-- +-- where \( Q \in \mathbb{R}^{m \times m} \) is orthogonal and +-- \( R \in \mathbb{R}^{m \times n} \) is upper triangular. +-- +-- The implementation uses in-place mutable arrays: the Householder +-- vectors are stored in the subdiagonal of the working copy of \( A \) +-- (LAPACK convention), and \( Q \) is formed via backward accumulation. +-- Total allocation: two matrices (one for \( R \), one for \( Q \)), +-- plus a small vector of \( \beta \) scalars. +-- +-- __Complexity:__ \( 2mn^2 - \tfrac{2}{3}n^3 \) flops for the +-- triangularisation, plus \( 2m^2 n - \tfrac{2}{3}n^3 \) flops for +-- accumulating \( Q \) (GVL4 p. 249). +-- +-- See also 'qrR' when only \( R \) is needed, and 'qrGivens' for a +-- Givens-based alternative. +qr :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix m n r e -> (Matrix m m r e, Matrix m n r e) +qr a = + let mm = dimVal @m + nn = dimVal @n + steps = min mm nn + + -- Phase 1: In-place Householder triangularisation. + -- After this, rArr holds R in the upper triangle and Householder + -- vectors v_k in the subdiagonal of column k (with v_k(k) = 1 implicit). + -- betaList holds the β scalars as a Haskell list. + (betaList, rArr) = M.withMArrayST (unMatrix a) $ \mr -> do + betas <- mapM (\k -> do + -- Compute Householder vector for column k, rows k..m-1 + x0 <- M.readM mr (k :. k) + sigma <- sumSqRange mr k mm k -- σ = Σ R(i,k)² for i=k+1..m-1 + if sigma == 0 && x0 >= 0 + then pure 0 -- already in desired form + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + -- Store v_k in subdiagonal of column k (v_k(k) = 1 is implicit) + -- Scale entries: v(i) = R(i,k) / v0 for i > k + mapM_ (\i -> do + rik <- M.readM mr (i :. k) + M.write_ mr (i :. k) (rik / v0) + ) [k+1..mm-1] + -- Set diagonal: R(k,k) = μ (result of H_k applied to column k) + M.write_ mr (k :. k) mu + -- Apply H_k from the left to columns k+1..n-1 of R: + -- (Column k is skipped: its diagonal is μ, subdiagonal stores v) + -- For each column j: w_j = β(R(k,j) + Σ_{i>k} v(i)·R(i,j)) + -- R(k,j) -= w_j; R(i,j) -= v(i)·w_j + mapM_ (\j -> do + rkj <- M.readM mr (k :. j) + wj <- sumProdRange mr mr k mm k j + let wj' = beta * (rkj + wj) + M.write_ mr (k :. j) (rkj - wj') + mapM_ (\i -> do + vi <- M.readM mr (i :. k) + rij <- M.readM mr (i :. j) + M.write_ mr (i :. j) (rij - vi * wj') + ) [k+1..mm-1] + ) [k+1..nn-1] + pure beta + ) [0..steps-1] + pure betas + + -- Phase 2: Backward accumulation of Q. + -- Start with Q = I, then for k = steps-1 downto 0: + -- Apply H_k from the right: Q <- Q·(I - β_k·v_k·v_k^T) + qMat = createMatrix @m @m @r $ \mq -> do + -- Initialize Q = I + mapM_ (\i -> mapM_ (\j -> + M.write_ mq (i :. j) (if i == j then 1 else 0) + ) [0..mm-1]) [0..mm-1] + -- Forward accumulation: Q <- Q·H_0·H_1·…·H_{n-1} + mapM_ (\k -> do + let beta_k = betaList !! k + if beta_k == 0 then pure () + else + -- Apply (I - β·v·v^T) from the right to Q + -- For each row i: w_i = β·(Q(i,k) + Σ_{l>k} Q(i,l)·v(l)) + -- Q(i,k) -= w_i; Q(i,l) -= w_i·v(l) + mapM_ (\i -> do + qik <- M.readM mq (i :. k) + wi <- qvProd mq rArr i k mm + let wi' = beta_k * (qik + wi) + M.write_ mq (i :. k) (qik - wi') + mapM_ (\l -> do + let vl = M.index' rArr (l :. k) + qil <- M.readM mq (i :. l) + M.write_ mq (i :. l) (qil - wi' * vl) + ) [k+1..mm-1] + ) [0..mm-1] + ) [0..steps-1] + + -- Extract clean R (zero out subdiagonal, which holds Householder vectors) + rClean = makeMatrix @m @n @r $ \i j -> + if i <= j then M.index' rArr (i :. j) else 0 + + in (qMat, rClean) + +-- | Specialised QR factorisation for @P Double@ using raw ByteArray# primops. +qrP :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> (Matrix m m M.P Double, Matrix m n M.P Double) +qrP a = + let mm = dimVal @m + nn = dimVal @n + steps = min mm nn + + -- Phase 1: In-place Householder triangularisation using raw kernels. + (betaList, rArr) = M.withMArrayST (unMatrix a) $ \mr -> do + let mbaR = unwrapMutableByteArray mr + offR = unwrapMutableByteArrayOffset mr + betas <- mapM (\k -> do + -- Read x0 = R[k,k] + x0 <- M.readM mr (k :. k) + -- σ = Σ R[i,k]² for i = k+1..m-1 (using raw kernel) + sigma <- rawMutSumSqColumn mbaR offR nn (k + 1) mm k + if sigma == 0 && x0 >= 0 + then pure 0 + else do + let mu = sqrt (x0 * x0 + sigma) + v0 = if x0 <= 0 then x0 - mu else -sigma / (x0 + mu) + beta = 2 * v0 * v0 / (sigma + v0 * v0) + -- Scale v: v(i) = R(i,k) / v0 for i > k + scaleColumn mbaR offR nn (k + 1) mm k (1.0 / v0) + -- Set diagonal + M.write_ mr (k :. k) mu + -- Apply H_k to columns k+1..n-1 using raw kernel + mapM_ (\j -> + rawMutHouseholderApply mbaR offR nn beta k mm j + ) [k+1..nn-1] + pure beta + ) [0..steps-1] + pure betas + + -- Phase 2: Blocked WY Q accumulation using GEMM. + -- Group Householder vectors into panels of size nb, compute T-factor + -- via GEMM, and apply block reflectors via Level-3 GEMM. + qMat = createMatrix @m @m @M.P $ \mq -> do + let mbaQ = unwrapMutableByteArray mq + offQ = unwrapMutableByteArrayOffset mq + baR = unwrapByteArray rArr + offFR = unwrapByteArrayOffset rArr + -- Initialize Q = I via rawZeroDoubles + diagonal writes + rawZeroDoubles mbaQ offQ (mm * mm) + mapM_ (\i -> writeRawD mbaQ offQ (i * mm + i) 1) [0..mm-1] + + if steps <= 16 + then + -- Small matrix: per-row accumulation (Level-2) + mapM_ (\k -> do + let beta_k = betaList !! k + if beta_k == 0 then pure () + else + mapM_ (\i -> + rawMutQAccum mbaQ offQ mm baR offFR nn beta_k k mm i + ) [0..mm-1] + ) [0..steps-1] + else do + -- Blocked WY: batch nb Householder vectors at a time + let !nb = min 32 steps + mbaBetas <- newByteArray (steps * 8) + mapM_ (\(i, b) -> writeRawD mbaBetas 0 i b) (zip [0..] betaList) + mbaY <- newByteArray (mm * nb * 8) + mbaTf <- newByteArray (nb * nb * 8) + mbaW1 <- newByteArray (mm * nb * 8) + mbaW2 <- newByteArray (mm * nb * 8) + mbaYT <- newByteArray (nb * mm * 8) + mbaG <- newByteArray (nb * nb * 8) + + let goBlock !k0 + | k0 >= steps = pure () + | otherwise = do + let !bs = min nb (steps - k0) + -- Pack Y (mm × bs): Y[:,j] = v_{k0+j} + -- v_k has implicit 1 at position k, stored values at k+1..mm-1 + rawZeroDoubles mbaY 0 (mm * bs) + mapM_ (\j -> do + let !k = k0 + j + writeRawD mbaY 0 (k * bs + j) 1.0 + mapM_ (\l -> + writeRawD mbaY 0 (l * bs + j) (indexRawD baR offFR (l * nn + k)) + ) [k+1..mm-1] + ) [0..bs-1] + + -- Transpose Y → Y^T (bs × mm) for GEMM reuse + rawZeroDoubles mbaYT 0 (bs * mm) + mapM_ (\j -> do + let !k = k0 + j + writeRawD mbaYT 0 (j * mm + k) 1.0 + mapM_ (\l -> + writeRawD mbaYT 0 (j * mm + l) (indexRawD baR offFR (l * nn + k)) + ) [k+1..mm-1] + ) [0..bs-1] + + baY <- unsafeFreezeByteArray mbaY + baYT <- unsafeFreezeByteArray mbaYT + + -- Compute G = Y^T × Y (bs × bs) via GEMM + rawZeroDoubles mbaG 0 (bs * bs) + rawGemmKernel baYT 0 baY 0 mbaG 0 bs mm bs + + -- Build T-factor (bs × bs upper-triangular) + rawZeroDoubles mbaTf 0 (bs * bs) + mapM_ (\j -> do + betaj <- readRawD mbaBetas 0 (k0 + j) + writeRawD mbaTf 0 (j * bs + j) betaj + when (j > 0 && betaj /= 0) $ do + -- T[0..j-1, j] = -betaj * T[0..j-1, 0..j-1] * G[0..j-1, j] + mapM_ (\i -> do + g_ij <- readRawD mbaG 0 (i * bs + j) + writeRawD mbaW1 0 i g_ij + ) [0..j-1] + mapM_ (\i -> do + let triLoop !l !acc + | l >= j = pure acc + | otherwise = do + til <- readRawD mbaTf 0 (i * bs + l) + dl <- readRawD mbaW1 0 l + triLoop (l+1) (acc + til * dl) + z <- triLoop 0 0 + writeRawD mbaTf 0 (i * bs + j) (negate betaj * z) + ) [0..j-1] + ) [0..bs-1] + + -- W1 = Q · Y (mm×mm * mm×bs → mm×bs) + baQ <- unsafeFreezeByteArray mbaQ + rawZeroDoubles mbaW1 0 (mm * bs) + rawGemmKernel baQ offQ baY 0 mbaW1 0 mm mm bs + + -- W2 = W1 · T (mm×bs * bs×bs → mm×bs) + baW1 <- unsafeFreezeByteArray mbaW1 + baTf <- unsafeFreezeByteArray mbaTf + rawZeroDoubles mbaW2 0 (mm * bs) + rawGemmKernel baW1 0 baTf 0 mbaW2 0 mm bs bs + + -- Negate W2 in-place + rawNegateDoubles mbaW2 0 (mm * bs) + + -- Q += (-W2) · Y^T (mm×bs * bs×mm → mm×mm) + baNW2 <- unsafeFreezeByteArray mbaW2 + rawGemmKernel baNW2 0 baYT 0 mbaQ offQ mm bs mm + + goBlock (k0 + bs) + goBlock 0 + + -- Extract clean R + rClean = makeMatrix @m @n @M.P $ \i j -> + if i <= j then M.index' rArr (i :. j) else 0 + + in (qMat, rClean) +{-# NOINLINE qrP #-} + +-- | Scale elements in a column of a mutable matrix: A[i,col] *= scale for i in [start..end-1]. +scaleColumn :: MutableByteArray s -> Int -> Int -> Int -> Int -> Int -> Double -> ST s () +scaleColumn (MutableByteArray mba) (I# off) (I# ncols) (I# start) (I# end) (I# col) (D# scale) = ST $ \s0 -> + let go i s + | isTrue# (i >=# end) = s + | otherwise = + case readDoubleArray# mba (off +# i *# ncols +# col) s of + (# s', v #) -> + case writeDoubleArray# mba (off +# i *# ncols +# col) (v *## scale) s' of + s'' -> go (i +# 1#) s'' + in (# go start s0, () #) +{-# INLINE scaleColumn #-} + +-- | Helper: Σ R(i,col)² for i=start+1..end-1 (sum of squares below diagonal) +sumSqRange :: (M.Manifest r e, Num e) => M.MArray s r Ix2 e -> Int -> Int -> Int -> ST s e +sumSqRange mr start end col = go (start + 1) 0 + where + go i !acc + | i >= end = pure acc + | otherwise = do + v <- M.readM mr (i :. col) + go (i + 1) (acc + v * v) + +-- | Helper: Σ mr1(i,c1)·mr2(i,c2) for i=start+1..end-1 +sumProdRange :: (M.Manifest r e, Num e) + => M.MArray s r Ix2 e -> M.MArray s r Ix2 e + -> Int -> Int -> Int -> Int -> ST s e +sumProdRange mr1 mr2 start end c1 c2 = go (start + 1) 0 + where + go i !acc + | i >= end = pure acc + | otherwise = do + v1 <- M.readM mr1 (i :. c1) + v2 <- M.readM mr2 (i :. c2) + go (i + 1) (acc + v1 * v2) + +-- | Helper: Σ Q(row,l)·v(l) for l=start+1..end-1 where v is stored in rArr subdiag of col start +qvProd :: (M.Manifest r1 e, M.Manifest r2 e, Num e) + => M.MArray s r1 Ix2 e -> M.Array r2 Ix2 e -> Int -> Int -> Int -> ST s e +qvProd mq rArr row start end = go (start + 1) 0 + where + go l !acc + | l >= end = pure acc + | otherwise = do + qrl <- M.readM mq (row :. l) + let vl = M.index' rArr (l :. start) + go (l + 1) (acc + qrl * vl) + +-- | Compute only the upper triangular factor \( R \) from the QR +-- factorisation, without explicitly forming \( Q \) +-- (GVL4 Algorithm 5.2.1, p. 249). +-- +-- Unlike 'qr', this function never forms the orthogonal factor, +-- saving \( O(m^2 n) \) flops. The Householder vectors are computed +-- and applied in-place but discarded. +-- +-- __Complexity:__ \( 2mn^2 - \tfrac{2}{3}n^3 \) flops. +qrR :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix m n r e -> Matrix m n r e +qrR a = snd (qr a) + +-- | QR factorisation via Givens rotations (GVL4 Section 5.2.4, p. 252). +-- +-- Given \( A \in \mathbb{R}^{m \times n} \) with \( m \ge n \), compute +-- the factorisation \( A = Q \, R \) by applying a sequence of Givens +-- rotations to zero out sub-diagonal entries one at a time. +-- +-- The implementation computes \( R \) in-place via the 'ST' monad, +-- records the rotation parameters, then applies them to form \( Q \) +-- in a second in-place pass. +-- +-- __Complexity:__ \( 3mn^2 - n^3 \) flops for a dense \( m \times n \) +-- matrix; \( O(mn) \) flops for an upper Hessenberg matrix +-- (GVL4 p. 253). +qrGivens :: forall m n r e. (KnownNat m, KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix m n r e -> (Matrix m m r e, Matrix m n r e) +qrGivens a = + let mm = dimVal @m + nn = dimVal @n + steps = min mm nn + + -- Pass 1: compute R in-place, recording Givens rotations + (rots, rArr) = M.withMArrayST (unMatrix a) $ \mr -> do + let go j !acc + | j >= steps = pure acc + | otherwise = do + acc' <- goRows j (j+1) acc mr + go (j+1) acc' + goRows j i !acc mr_ + | i >= mm = pure acc + | otherwise = do + aij <- M.readM mr_ (i :. j) + if aij == 0 then goRows j (i+1) acc mr_ + else do + ajj <- M.readM mr_ (j :. j) + let (c, s) = givensRotation ajj aij + -- Apply G^T to rows j and i + mapM_ (\col -> do + rjc <- M.readM mr_ (j :. col) + ric <- M.readM mr_ (i :. col) + M.write_ mr_ (j :. col) (c * rjc - s * ric) + M.write_ mr_ (i :. col) (s * rjc + c * ric) + ) [0..nn-1] + goRows j (i+1) (acc ++ [(c, s, j, i)]) mr_ + go 0 [] + + -- Pass 2: form Q by applying recorded rotations to I + qMat = createMatrix @m @m @r $ \mq -> do + -- Initialize Q = I + mapM_ (\i -> mapM_ (\j -> + M.write_ mq (i :. j) (if i == j then 1 else 0) + ) [0..mm-1]) [0..mm-1] + -- Apply each rotation from the right: Q <- Q·G + mapM_ (\(c, s, ci, ck) -> + mapM_ (\row -> do + qrc <- M.readM mq (row :. ci) + qrk <- M.readM mq (row :. ck) + M.write_ mq (row :. ci) (c * qrc - s * qrk) + M.write_ mq (row :. ck) (s * qrc + c * qrk) + ) [0..mm-1] + ) rots + + in (qMat, MkMatrix rArr) + +-- Raw ByteArray# helpers for blocked WY Q accumulation +readRawD :: MutableByteArray s -> Int -> Int -> ST s Double +readRawD (MutableByteArray mba) (I# off) (I# i) = ST $ \s -> + case readDoubleArray# mba (off +# i) s of + (# s', v #) -> (# s', D# v #) +{-# INLINE readRawD #-} + +writeRawD :: MutableByteArray s -> Int -> Int -> Double -> ST s () +writeRawD (MutableByteArray mba) (I# off) (I# i) (D# v) = ST $ \s -> + case writeDoubleArray# mba (off +# i) v s of + s' -> (# s', () #) +{-# INLINE writeRawD #-} + +indexRawD :: ByteArray -> Int -> Int -> Double +indexRawD (ByteArray ba) (I# off) (I# i) = + D# (indexDoubleArray# ba (off +# i)) +{-# INLINE indexRawD #-} diff --git a/src/Numeric/LinearAlgebra/Massiv/Solve/Banded.hs b/src/Numeric/LinearAlgebra/Massiv/Solve/Banded.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Solve/Banded.hs @@ -0,0 +1,342 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Solve.Banded +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = Banded and Tridiagonal System Solvers +-- +-- Specialised LU and Cholesky factorizations for banded matrices, plus a +-- dedicated tridiagonal solver, following Golub & Van Loan, +-- /Matrix Computations/, 4th edition (GVL4), Section 4.3, pp. 174--182. +-- +-- A matrix \(A \in \mathbb{R}^{n \times n}\) has /lower bandwidth/ \(p\) +-- and /upper bandwidth/ \(q\) when \(a_{ij} = 0\) for \(i > j + p\) or +-- \(j > i + q\). Exploiting this structure reduces the factorization cost +-- from \(O(n^3)\) to \(O(npq)\) (GVL4 p. 176), and the triangular-solve +-- cost from \(O(n^2)\) to \(O(np)\) or \(O(nq)\). The important special +-- case of a tridiagonal system (\(p = q = 1\)) is solvable in \(O(n)\) +-- flops. +-- +-- +-------------------+-----------------------------------+----------------------------------+ +-- | Function | Algorithm | Reference | +-- +===================+===================================+==================================+ +-- | 'bandLU' | Band Gaussian elimination | GVL4 Algorithm 4.3.1, p. 175 | +-- +-------------------+-----------------------------------+----------------------------------+ +-- | 'bandForwardSub' | Band forward substitution | GVL4 Algorithm 4.3.2, p. 176 | +-- +-------------------+-----------------------------------+----------------------------------+ +-- | 'bandBackSub' | Band back substitution | GVL4 Algorithm 4.3.3, p. 176 | +-- +-------------------+-----------------------------------+----------------------------------+ +-- | 'bandCholesky' | Band Cholesky factorization | GVL4 Algorithm 4.3.5, p. 178 | +-- +-------------------+-----------------------------------+----------------------------------+ +-- | 'tridiagSolve' | SPD tridiagonal solver | GVL4 Algorithm 4.3.6, p. 179 | +-- +-------------------+-----------------------------------+----------------------------------+ +-- +-- == Complexity +-- +-- * Band LU: \(O(npq)\) flops (GVL4 p. 176). +-- * Band triangular solve: \(O(np)\) or \(O(nq)\) flops (GVL4 p. 176). +-- * Band Cholesky: \(O(np^2)\) flops (GVL4 p. 178). +-- * Tridiagonal solver: \(O(n)\) flops (GVL4 p. 179). +-- +-- == Type Safety +-- +-- Matrix dimensions are tracked at the type level via 'KnownNat'. The +-- bandwidths \(p\) and \(q\) are passed as run-time 'Int' values because +-- they are often data-dependent. The matrix is stored in standard dense +-- format; only the band is accessed. +module Numeric.LinearAlgebra.Massiv.Solve.Banded + ( -- * Band LU factorization + bandLU + -- * Band triangular solve + , bandForwardSub + , bandBackSub + -- * Band Cholesky (\(A = GG^T\)) + , bandCholesky + -- * Tridiagonal solver (\(Ax = b\), \(p = q = 1\)) + , tridiagSolve + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal + +-- | Band Gaussian elimination without pivoting (GVL4 Algorithm 4.3.1, +-- p. 175). +-- +-- Given an \(n \times n\) matrix \(A\) with lower bandwidth \(p\) and upper +-- bandwidth \(q\), computes an in-place \(LU\) factorization where \(L\) +-- has lower bandwidth \(p\) and \(U\) has upper bandwidth \(q\). +-- +-- The matrix is stored in standard dense format; only the entries within +-- the band are accessed or modified. +-- +-- __Precondition.__ All leading principal submatrices of \(A\) must be +-- nonsingular (analogous to the dense case, GVL4 Theorem 3.2.1, p. 116). +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) statically ensures the matrix is square. The bandwidth +-- parameters \(p\) and \(q\) are run-time values. +-- +-- ==== Complexity +-- +-- \(O(npq)\) flops (GVL4 p. 176). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 4.3.1 +-- (Band Gaussian Elimination), p. 175. +bandLU :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e) + => Int -- ^ Lower bandwidth \(p\) + -> Int -- ^ Upper bandwidth \(q\) + -> Matrix n n r e -> Matrix n n r e +bandLU p q (MkMatrix a) = + let nn = dimVal @n + in MkMatrix $ snd $ M.withMArrayST a $ \ma -> + mapM_ (\k -> do + akk <- M.readM ma (k :. k) + let imax = min (k + p) (nn - 1) + -- Compute multipliers + mapM_ (\i -> do + aik <- M.readM ma (i :. k) + M.write_ ma (i :. k) (aik / akk) + ) [k+1..imax] + -- Update + let jmax = min (k + q) (nn - 1) + mapM_ (\j -> + mapM_ (\i -> do + aij <- M.readM ma (i :. j) + aik <- M.readM ma (i :. k) + akj <- M.readM ma (k :. j) + M.write_ ma (i :. j) (aij - aik * akj) + ) [k+1..imax] + ) [k+1..jmax] + ) [0..nn-2] + +-- | Band forward substitution (GVL4 Algorithm 4.3.2, p. 176). +-- +-- Solves \(Lx = b\) where \(L\) is /unit/ lower triangular with lower +-- bandwidth \(p\). Only the \(p\) subdiagonals are accessed; the unit +-- diagonal is implicit, so no division is performed and the constraint +-- relaxes to 'Num'. +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) ensures the dimensions of \(L\) and \(b\) agree at +-- compile time. +-- +-- ==== Complexity +-- +-- \(O(np)\) flops (GVL4 p. 176). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 4.3.2 +-- (Band Forward Substitution), p. 176. +bandForwardSub :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Int -- ^ Lower bandwidth \(p\) + -> Matrix n n r e -> Vector n r e -> Vector n r e +bandForwardSub p l b = createVector @n $ \mx -> do + let nn = dimVal @n + mapM_ (\i -> M.write_ mx i (b !. i)) [0..nn-1] + mapM_ (\j -> do + xj <- M.readM mx j + let imax = min (j + p) (nn - 1) + mapM_ (\i -> do + xi <- M.readM mx i + M.write_ mx i (xi - (l ! (i, j)) * xj) + ) [j+1..imax] + ) [0..nn-1] + +-- | Band back substitution (GVL4 Algorithm 4.3.3, p. 176). +-- +-- Solves \(Ux = b\) where \(U\) is upper triangular with upper bandwidth +-- \(q\). Only the diagonal and the \(q\) superdiagonals are accessed. +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) ensures the dimensions of \(U\) and \(b\) agree at +-- compile time. +-- +-- ==== Complexity +-- +-- \(O(nq)\) flops (GVL4 p. 176). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 4.3.3 +-- (Band Back Substitution), p. 176. +bandBackSub :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e) + => Int -- ^ Upper bandwidth \(q\) + -> Matrix n n r e -> Vector n r e -> Vector n r e +bandBackSub q u b = createVector @n $ \mx -> do + let nn = dimVal @n + mapM_ (\i -> M.write_ mx i (b !. i)) [0..nn-1] + mapM_ (\j -> do + xj <- M.readM mx j + let ujj = u ! (j, j) + xj' = xj / ujj + M.write_ mx j xj' + let imin = max 0 (j - q) + mapM_ (\i -> do + xi <- M.readM mx i + M.write_ mx i (xi - (u ! (i, j)) * xj') + ) [imin..j-1] + ) [nn-1, nn-2..0] + +-- | Band Cholesky factorization (GVL4 Algorithm 4.3.5, p. 178). +-- +-- Given a symmetric positive definite \(n \times n\) matrix \(A\) with +-- lower bandwidth \(p\), computes the lower triangular factor \(G\) (also +-- with bandwidth \(p\)) such that +-- +-- \[ +-- A = G G^T +-- \] +-- +-- Only the lower band of \(A\) (entries \(a_{ij}\) with +-- \(0 \le i - j \le p\)) is read. +-- +-- __Precondition.__ \(A\) must be symmetric positive definite. If this +-- condition is violated, the algorithm may encounter a negative value under +-- a square root and produce NaN. +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) statically ensures the matrix is square. The bandwidth +-- \(p\) is a run-time value. +-- +-- ==== Complexity +-- +-- \(O(np^2)\) flops (GVL4 p. 178). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 4.3.5 +-- (Band Cholesky), p. 178. +bandCholesky :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Int -- ^ Bandwidth \(p\) + -> Matrix n n r e -> Matrix n n r e +bandCholesky p (MkMatrix a) = + let nn = dimVal @n + in MkMatrix $ M.createArrayST_ (M.Sz2 nn nn) $ \mg -> do + -- Initialize to zero + mapM_ (\i -> mapM_ (\j -> M.write_ mg (i :. j) 0) [0..nn-1]) [0..nn-1] + -- Copy lower band of A + mapM_ (\j -> + let imax = min (j + p) (nn - 1) + in mapM_ (\i -> M.write_ mg (i :. j) (M.index' a (i :. j))) [j..imax] + ) [0..nn-1] + -- Band Cholesky + mapM_ (\j -> do + -- Subtract contributions + let kmin = max 0 (j - p) + mapM_ (\k -> do + gjk <- M.readM mg (j :. k) + let lam = min (j + p) (nn - 1) + mapM_ (\i -> do + gij <- M.readM mg (i :. j) + gik <- M.readM mg (i :. k) + M.write_ mg (i :. j) (gij - gik * gjk) + ) [j..lam] + ) [kmin..j-1] + -- Scale + gjj <- M.readM mg (j :. j) + let sjj = sqrt gjj + lam = min (j + p) (nn - 1) + mapM_ (\i -> do + gij <- M.readM mg (i :. j) + M.write_ mg (i :. j) (gij / sjj) + ) [j..lam] + ) [0..nn-1] + +-- | Symmetric positive definite tridiagonal system solver (GVL4 +-- Algorithm 4.3.6, p. 179). +-- +-- Solves \(Ax = b\) where \(A\) is symmetric, tridiagonal, and positive +-- definite. The matrix \(A\) is specified compactly by its diagonal +-- \(\alpha_{1:n}\) and its superdiagonal \(\beta_{1:n-1}\). +-- +-- The algorithm computes the \(LDL^T\) factorization of \(A\) and folds it +-- together with forward and back substitution in a single \(O(n)\) pass: +-- +-- 1. __Factor:__ Compute \(A = LDL^T\) where \(L\) is unit lower +-- bidiagonal and \(D\) is diagonal. +-- 2. __Forward substitution:__ Solve \(Lz = b\). +-- 3. __Diagonal solve:__ Solve \(Dy = z\). +-- 4. __Back substitution:__ Solve \(L^T x = y\). +-- +-- __Precondition.__ \(A\) must be symmetric positive definite. This is not +-- checked at run time. +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) enforces that the diagonal, superdiagonal, right-hand +-- side, and solution vectors all have length \(n\) at compile time. +-- +-- ==== Complexity +-- +-- \(O(n)\) flops (GVL4 p. 179). This is optimal for tridiagonal systems. +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 4.3.6 +-- (SPD Tridiagonal System Solver), p. 179. +tridiagSolve :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e) + => Vector n r e -- ^ Diagonal entries \(\alpha_{1:n}\) + -> Vector n r e -- ^ Superdiagonal entries \(\beta_{1:n-1}\) (length \(n\), only indices @0..n-2@ used) + -> Vector n r e -- ^ Right-hand side \(b\) + -> Vector n r e -- ^ Solution \(x\) +tridiagSolve diag supdiag b = createVector @n $ \mx -> do + let nn = dimVal @n + -- Working arrays for modified diagonal and superdiagonal + alpha <- M.newMArray @r (M.Sz1 nn) (0 :: e) + beta <- M.newMArray @r (M.Sz1 nn) (0 :: e) + -- Initialize + mapM_ (\i -> do + M.write_ alpha i (diag !. i) + M.write_ beta i (if i < nn - 1 then supdiag !. i else 0) + ) [0..nn-1] + -- Copy b into result + mapM_ (\i -> M.write_ mx i (b !. i)) [0..nn-1] + + -- LDLᵀ factorization and forward substitution combined + -- for k = 2:n + -- t = β(k-1), β(k-1) = t/α(k-1), α(k) = α(k) - t·β(k-1) + mapM_ (\k -> do + t <- M.readM beta (k - 1) + ak1 <- M.readM alpha (k - 1) + let bk1 = t / ak1 + M.write_ beta (k - 1) bk1 + ak <- M.readM alpha k + M.write_ alpha k (ak - t * bk1) + ) [1..nn-1] + + -- Forward substitution: b(k) = b(k) - β(k-1)·b(k-1) + mapM_ (\k -> do + bk <- M.readM mx k + bk1 <- M.readM mx (k - 1) + betaK1 <- M.readM beta (k - 1) + M.write_ mx k (bk - betaK1 * bk1) + ) [1..nn-1] + + -- Diagonal solve: b(n) = b(n)/α(n) + bn <- M.readM mx (nn - 1) + an <- M.readM alpha (nn - 1) + M.write_ mx (nn - 1) (bn / an) + + -- Back substitution: b(k) = b(k)/α(k) - β(k)·b(k+1) + mapM_ (\k -> do + bk <- M.readM mx k + ak <- M.readM alpha k + betaK <- M.readM beta k + bk1 <- M.readM mx (k + 1) + M.write_ mx k (bk / ak - betaK * bk1) + ) [nn-2, nn-3..0] diff --git a/src/Numeric/LinearAlgebra/Massiv/Solve/Cholesky.hs b/src/Numeric/LinearAlgebra/Massiv/Solve/Cholesky.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Solve/Cholesky.hs @@ -0,0 +1,449 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Solve.Cholesky +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = Cholesky Factorization +-- +-- Cholesky decomposition for symmetric positive definite (SPD) matrices, +-- following Golub & Van Loan, /Matrix Computations/, 4th edition (GVL4), +-- Section 4.2, pp. 163--169. +-- +-- For any SPD matrix \(A \in \mathbb{R}^{n \times n}\), there exists a +-- unique lower triangular matrix \(G\) with positive diagonal entries such +-- that +-- +-- \[ +-- A = G G^T +-- \] +-- +-- (GVL4 Theorem 4.2.1, p. 163). This is the /Cholesky factorization/. +-- Because it exploits symmetry, the Cholesky factorization requires only +-- half the work of a general LU factorization: \(O(n^3/3)\) flops vs. +-- \(O(2n^3/3)\) flops (GVL4 p. 165). +-- +-- +-------------------+-------------------------------+---------------------------------+ +-- | Function | Algorithm | Reference | +-- +===================+===============================+=================================+ +-- | 'cholesky' | Outer-product Cholesky | GVL4 Algorithm 4.2.1, p. 164 | +-- +-------------------+-------------------------------+---------------------------------+ +-- | 'choleskyGaxpy' | Gaxpy (column-oriented) | GVL4 Algorithm 4.2.2, p. 165 | +-- +-------------------+-------------------------------+---------------------------------+ +-- | 'choleskySolve' | Solve via \(A = GG^T\) | GVL4 Section 4.2, p. 166 | +-- +-------------------+-------------------------------+---------------------------------+ +-- +-- == Complexity +-- +-- The factorization costs \(O(n^3/3)\) flops -- exactly half of LU +-- (GVL4 p. 165). The subsequent pair of triangular solves adds \(O(n^2)\) +-- flops. +-- +-- == Type Safety +-- +-- Matrix dimensions are tracked at the type level via 'KnownNat', so the +-- compiler statically ensures the coefficient matrix is square and the +-- right-hand side vector has a conforming length. Note that positive +-- definiteness is /not/ checked at the type level; if the input matrix is +-- not SPD, the algorithm may produce NaN values from taking the square root +-- of a negative number. +module Numeric.LinearAlgebra.Massiv.Solve.Cholesky + ( -- * Cholesky factorization (\(A = GG^T\)) + cholesky + , choleskyGaxpy + -- * Solving with Cholesky (\(Ax = b\)) + , choleskySolve + , choleskySolveP + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..), unwrapByteArray, unwrapByteArrayOffset, + unwrapMutableByteArray, unwrapMutableByteArrayOffset) +import GHC.TypeNats (KnownNat) +import Control.Monad (when) +import GHC.Exts +import GHC.Prim +import GHC.ST (ST(..)) +import Data.Primitive.ByteArray (ByteArray(..), MutableByteArray(..), newByteArray, unsafeFreezeByteArray) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Solve.Triangular (forwardSub, backSub) +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (transpose) +import Numeric.LinearAlgebra.Massiv.Internal.Kernel + (rawCholColumnSIMD, rawCholColumnSIMDFrom, + rawForwardSubCholPackedSIMD, rawBackSubCholTPackedSIMD, + rawGemmKernel, rawZeroDoubles) + +-- | Outer-product Cholesky factorization (GVL4 Algorithm 4.2.1, p. 164). +-- +-- Given a symmetric positive definite \(n \times n\) matrix \(A\), computes +-- the unique lower triangular matrix \(G\) with positive diagonal entries +-- such that +-- +-- \[ +-- A = G G^T +-- \] +-- +-- Only the lower triangle of \(A\) is accessed; the upper triangle is +-- ignored. +-- +-- The algorithm processes one column at a time using an /outer-product/ +-- update of the trailing submatrix. +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) statically ensures \(A\) is square. The 'Floating' +-- constraint provides 'sqrt'. Positive definiteness is a run-time +-- precondition; violation may produce NaN from \(\sqrt{g_{jj}}\) when +-- \(g_{jj} < 0\). +-- +-- ==== Complexity +-- +-- \(O(n^3/3)\) flops (GVL4 p. 165). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 4.2.1 +-- (Outer Product Cholesky), p. 164. Existence: Theorem 4.2.1, p. 163. +cholesky :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Matrix n n r e +cholesky (MkMatrix a) = + let nn = dimVal @n + in MkMatrix $ M.createArrayST_ (M.Sz2 nn nn) $ \mg -> do + -- Initialize to zero + mapM_ (\i -> mapM_ (\j -> M.write_ mg (i :. j) 0) [0..nn-1]) [0..nn-1] + -- Copy lower triangle of A into working storage + mapM_ (\j -> mapM_ (\i -> do + let aij = M.index' a (i :. j) + M.write_ mg (i :. j) aij + ) [j..nn-1]) [0..nn-1] + + -- Outer product Cholesky + mapM_ (\j -> do + -- Subtract contributions from previous columns + mapM_ (\k -> do + gjk <- M.readM mg (j :. k) + mapM_ (\i -> do + gij <- M.readM mg (i :. j) + gik <- M.readM mg (i :. k) + M.write_ mg (i :. j) (gij - gik * gjk) + ) [j..nn-1] + ) [0..j-1] + + -- Scale column + gjj <- M.readM mg (j :. j) + let sjj = sqrt gjj + M.write_ mg (j :. j) sjj + mapM_ (\i -> do + gij <- M.readM mg (i :. j) + M.write_ mg (i :. j) (gij / sjj) + ) [j+1..nn-1] + ) [0..nn-1] + +-- | Gaxpy (column-oriented) Cholesky factorization (GVL4 Algorithm 4.2.2, +-- p. 165). +-- +-- Functionally equivalent to 'cholesky', but uses a /gaxpy/ (generalised +-- @y <- y - Gx@) inner loop that accumulates updates into each column +-- before normalising. This access pattern is advantageous for column-major +-- storage because it streams through contiguous memory. +-- +-- ==== Mathematical definition +-- +-- Computes the same \(G\) such that \(A = GG^T\) as 'cholesky'. +-- +-- ==== Type-safety guarantees +-- +-- Identical to 'cholesky'. +-- +-- ==== Complexity +-- +-- \(O(n^3/3)\) flops (GVL4 p. 165). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 4.2.2 +-- (Gaxpy Cholesky), p. 165. +choleskyGaxpy :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Matrix n n r e +choleskyGaxpy (MkMatrix a) = + let nn = dimVal @n + in MkMatrix $ M.createArrayST_ (M.Sz2 nn nn) $ \mg -> do + -- Initialize: copy lower triangle of A + mapM_ (\i -> mapM_ (\j -> + if i >= j + then M.write_ mg (i :. j) (M.index' a (i :. j)) + else M.write_ mg (i :. j) 0 + ) [0..nn-1]) [0..nn-1] + + -- Column by column + mapM_ (\j -> do + -- Update column j using previous columns (gaxpy) + mapM_ (\k -> do + gjk <- M.readM mg (j :. k) + mapM_ (\i -> do + gij <- M.readM mg (i :. j) + gik <- M.readM mg (i :. k) + M.write_ mg (i :. j) (gij - gik * gjk) + ) [j..nn-1] + ) [0..j-1] + + -- Scale by 1/sqrt(g(j,j)) + gjj <- M.readM mg (j :. j) + let sjj = sqrt gjj + mapM_ (\i -> do + gij <- M.readM mg (i :. j) + M.write_ mg (i :. j) (gij / sjj) + ) [j..nn-1] + ) [0..nn-1] + +-- | Solve \(Ax = b\) where \(A\) is symmetric positive definite, using +-- Cholesky factorization (GVL4 Section 4.2, p. 166). +-- +-- The algorithm proceeds in three stages: +-- +-- 1. Factor \(A = GG^T\) via 'cholesky' (Algorithm 4.2.1). +-- 2. Solve \(Gy = b\) by forward substitution ('forwardSub'). +-- 3. Solve \(G^T x = y\) by back substitution ('backSub'). +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) enforces that \(A\) is \(n \times n\) and \(b\) has +-- length \(n\) at compile time. +-- +-- ==== Complexity +-- +-- \(O(n^3/3)\) flops for the factorization plus \(O(n^2)\) flops for the +-- two triangular solves (GVL4 p. 166). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Section 4.2, +-- pp. 163--169. +choleskySolve :: forall n r e. (KnownNat n, M.Manifest r e, Floating e, Ord e) + => Matrix n n r e -> Vector n r e -> Vector n r e +choleskySolve a b = + let g = cholesky a + gt = transpose g + y = forwardSub g b + in backSub gt y + +-- | Specialised Cholesky solve for @P Double@. +-- Does Cholesky factorisation + solve entirely using raw ByteArray# primops, +-- avoiding separate G/G^T matrix construction. +-- For n >= 64, uses panel (blocked) Cholesky factorisation with GEMM trailing update. +choleskySolveP :: forall n. KnownNat n + => Matrix n n M.P Double -> Vector n M.P Double -> Vector n M.P Double +choleskySolveP (MkMatrix a) (MkVector b) = + let nn = dimVal @n + in createVector @n @M.P $ \mx -> do + -- Allocate n×n working storage for G, copy lower triangle of A + mg <- M.newMArray @M.P (M.Sz2 nn nn) (0 :: Double) + let mbaG = unwrapMutableByteArray mg + offG = unwrapMutableByteArrayOffset mg + -- Copy lower triangle using raw primops + copyLowerTriangle a mbaG offG nn + + -- Phase 1: Cholesky factorisation + if nn >= 64 + then panelCholFactor mbaG offG nn 32 + else mapM_ (rawCholColumnSIMD mbaG offG nn) [0..nn-1] + + -- Phase 2: Freeze G and prepare RHS + frozenG <- M.freezeS mg + let baG = unwrapByteArray frozenG + offFG = unwrapByteArrayOffset frozenG + + -- Copy b into output vector + let mbaX = unwrapMutableByteArray mx + offX = unwrapMutableByteArrayOffset mx + copyVectorRaw b mbaX offX nn + + -- Phase 3: Forward substitution (Gy = b, SIMD dot-product) + rawForwardSubCholPackedSIMD baG offFG nn mbaX offX + + -- Phase 4: Back substitution (G^T x = y, SIMD SAXPY) + rawBackSubCholTPackedSIMD baG offFG nn mbaX offX +{-# NOINLINE choleskySolveP #-} + +-- | Panel (blocked) Cholesky factorisation with GEMM trailing update. +-- For each panel of width @nb@: +-- 1. Apply GEMM from previous panels: G[j:n, j:j+jb] -= L_prev × L_prev_panel^T +-- 2. Factor the panel using within-panel Cholesky (dot from panel start) +panelCholFactor :: MutableByteArray s -> Int -> Int -> Int -> ST s () +panelCholFactor mbaG offG nn nb = go 0 + where + go !j + | j >= nn = pure () + | otherwise = do + let !jb = min nb (nn - j) + !nBelow = nn - j -- rows in [j..n-1] + + -- Step 1: update current panel with contributions from previous panels + when (j > 0) $ do + -- L_below = G[j:n, 0:j], shape nBelow × j + -- L_panel = G[j:j+jb, 0:j], shape jb × j + -- Update: G[j:n, j:j+jb] -= L_below × L_panel^T + -- We compute this as: GEMM(L_below, L_panelT), where L_panelT = transpose(L_panel) + + -- Copy L_below to dense buffer (nBelow × j) + bufA <- newByteArray (nBelow * j * 8) + rawCopyCholSubmatrix mbaG offG nn j 0 nBelow j bufA 0 + + -- Copy L_panel transposed to dense buffer (j × jb) + -- L_panel is jb × j at rows [j..j+jb-1], cols [0..j-1] + -- Transposed: j × jb + bufBT <- newByteArray (j * jb * 8) + rawCopyCholTranspose mbaG offG nn j 0 jb j bufBT 0 + + -- Freeze for GEMM + baA <- unsafeFreezeByteArray bufA + baBT <- unsafeFreezeByteArray bufBT + + -- GEMM: C = L_below × L_panelT (nBelow × jb) + bufC <- newByteArray (nBelow * jb * 8) + rawZeroDoubles bufC 0 (nBelow * jb) + rawGemmKernel baA 0 baBT 0 bufC 0 nBelow j jb + + -- Subtract C from G[j:n, j:j+jb] + baC <- unsafeFreezeByteArray bufC + rawCholSubtractPanel baC 0 jb mbaG offG nn j j nBelow jb + + -- Step 2: factor panel using within-panel dependencies only + mapM_ (\c -> rawCholColumnSIMDFrom mbaG offG nn c j) [j..j+jb-1] + + go (j + jb) + +-- | Copy submatrix G[rowStart..rowStart+m-1, colStart..colStart+k-1] (stride n) +-- into dense buffer (stride k). +rawCopyCholSubmatrix :: MutableByteArray s -> Int -> Int + -> Int -> Int -> Int -> Int + -> MutableByteArray s -> Int + -> ST s () +rawCopyCholSubmatrix (MutableByteArray mba_src) (I# off_src) (I# n) + (I# rowStart) (I# colStart) (I# m) (I# k) + (MutableByteArray mba_dst) (I# off_dst) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# m) = s + | otherwise = + let srcRow = off_src +# (rowStart +# i) *# n +# colStart + dstRow = off_dst +# i *# k + span_ = k -# (k `remInt#` 4#) + goSimd j s_ + | isTrue# (j >=# span_) = s_ + | otherwise = + case readDoubleArrayAsDoubleX4# mba_src (srcRow +# j) s_ of + (# s1, v #) -> + case writeDoubleArrayAsDoubleX4# mba_dst (dstRow +# j) v s1 of + s2 -> goSimd (j +# 4#) s2 + goScalar j s_ + | isTrue# (j >=# k) = s_ + | otherwise = + case readDoubleArray# mba_src (srcRow +# j) s_ of + (# s1, v #) -> + case writeDoubleArray# mba_dst (dstRow +# j) v s1 of + s2 -> goScalar (j +# 1#) s2 + in goI (i +# 1#) (goScalar span_ (goSimd 0# s)) + in (# goI 0# s0, () #) +{-# INLINE rawCopyCholSubmatrix #-} + +-- | Copy and transpose: src[rowStart..rowStart+m-1, colStart..colStart+k-1] (stride n) +-- into dense buffer of shape k × m (stride m). +-- i.e., dst[j, i] = src[rowStart+i, colStart+j] +rawCopyCholTranspose :: MutableByteArray s -> Int -> Int + -> Int -> Int -> Int -> Int + -> MutableByteArray s -> Int + -> ST s () +rawCopyCholTranspose (MutableByteArray mba_src) (I# off_src) (I# n) + (I# rowStart) (I# colStart) (I# m) (I# k) + (MutableByteArray mba_dst) (I# off_dst) = ST $ \s0 -> + -- For each source row i, read k elements, write them as column i of dst + let goI i s + | isTrue# (i >=# m) = s + | otherwise = + let srcRow = off_src +# (rowStart +# i) *# n +# colStart + goJ j s_ + | isTrue# (j >=# k) = s_ + | otherwise = + case readDoubleArray# mba_src (srcRow +# j) s_ of + (# s1, v #) -> + -- dst[j, i] at offset j * m + i + case writeDoubleArray# mba_dst (off_dst +# j *# m +# i) v s1 of + s2 -> goJ (j +# 1#) s2 + in goI (i +# 1#) (goJ 0# s) + in (# goI 0# s0, () #) +{-# INLINE rawCopyCholTranspose #-} + +-- | Subtract dense buffer C (m × k, stride srcStride) from +-- G[rowStart..rowStart+m-1, colStart..colStart+k-1] (stride n). +rawCholSubtractPanel :: ByteArray -> Int -> Int + -> MutableByteArray s -> Int -> Int + -> Int -> Int -> Int -> Int + -> ST s () +rawCholSubtractPanel (ByteArray ba_src) (I# off_src) (I# srcStride) + (MutableByteArray mba_dst) (I# off_dst) (I# n) + (I# rowStart) (I# colStart) (I# m) (I# k) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# m) = s + | otherwise = + let srcRow = off_src +# i *# srcStride + dstRow = off_dst +# (rowStart +# i) *# n +# colStart + span_ = k -# (k `remInt#` 4#) + goSimd j s_ + | isTrue# (j >=# span_) = s_ + | otherwise = + case readDoubleArrayAsDoubleX4# mba_dst (dstRow +# j) s_ of + (# s1, aij #) -> + let cij = indexDoubleArrayAsDoubleX4# ba_src (srcRow +# j) + aij' = plusDoubleX4# aij (negateDoubleX4# cij) + in case writeDoubleArrayAsDoubleX4# mba_dst (dstRow +# j) aij' s1 of + s2 -> goSimd (j +# 4#) s2 + goScalar j s_ + | isTrue# (j >=# k) = s_ + | otherwise = + case readDoubleArray# mba_dst (dstRow +# j) s_ of + (# s1, aij #) -> + let cij = indexDoubleArray# ba_src (srcRow +# j) + in case writeDoubleArray# mba_dst (dstRow +# j) (aij -## cij) s1 of + s2 -> goScalar (j +# 1#) s2 + in goI (i +# 1#) (goScalar span_ (goSimd 0# s)) + in (# goI 0# s0, () #) +{-# INLINE rawCholSubtractPanel #-} + +-- | Copy lower triangle of an immutable 2D P array into a mutable byte array. +copyLowerTriangle :: M.Array M.P Ix2 Double -> MutableByteArray s -> Int -> Int -> ST s () +copyLowerTriangle src (MutableByteArray mba_dst) (I# off_dst) (I# n) = ST $ \s0 -> + let ba_src = case unwrapByteArray src of ByteArray ba -> ba + off_src = case unwrapByteArrayOffset src of I# o -> o + goI i s + | isTrue# (i >=# n) = s + | otherwise = goI (i +# 1#) (goJ i 0# s) + goJ i j s + | isTrue# (j ># i) = s + | otherwise = + let v = indexDoubleArray# ba_src (off_src +# i *# n +# j) + in case writeDoubleArray# mba_dst (off_dst +# i *# n +# j) v s of + s' -> goJ i (j +# 1#) s' + in (# goI 0# s0, () #) +{-# INLINE copyLowerTriangle #-} + +-- | Copy an immutable P vector into a mutable byte array. +copyVectorRaw :: M.Array M.P Ix1 Double -> MutableByteArray s -> Int -> Int -> ST s () +copyVectorRaw src (MutableByteArray mba_dst) (I# off_dst) (I# n) = ST $ \s0 -> + let ba_src = case unwrapByteArray src of ByteArray ba -> ba + off_src = case unwrapByteArrayOffset src of I# o -> o + go i s + | isTrue# (i >=# n) = s + | otherwise = + let v = indexDoubleArray# ba_src (off_src +# i) + in case writeDoubleArray# mba_dst (off_dst +# i) v s of + s' -> go (i +# 1#) s' + in (# go 0# s0, () #) +{-# INLINE copyVectorRaw #-} diff --git a/src/Numeric/LinearAlgebra/Massiv/Solve/LU.hs b/src/Numeric/LinearAlgebra/Massiv/Solve/LU.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Solve/LU.hs @@ -0,0 +1,562 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Solve.LU +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = LU Factorization +-- +-- LU decomposition with and without partial pivoting, plus derived +-- operations (linear solve and determinant), following Golub & Van Loan, +-- /Matrix Computations/, 4th edition (GVL4), Sections 3.2 and 3.4, +-- pp. 114--131. +-- +-- Given an \(n \times n\) matrix \(A\), the factorization produces +-- +-- \[ +-- PA = LU +-- \] +-- +-- where \(P\) is a permutation matrix, \(L\) is unit lower triangular, and +-- \(U\) is upper triangular (GVL4 Theorem 3.4.1, p. 125). +-- +-- Without pivoting (\(P = I\)) the factorization exists if and only if all +-- leading principal submatrices of \(A\) are nonsingular +-- (GVL4 Theorem 3.2.1, p. 116). Partial pivoting guarantees existence for +-- any nonsingular \(A\) and improves numerical stability by bounding the +-- growth factor (GVL4 Section 3.4.6). +-- +-- +-------------------+----------------------------------+-------------------------------+ +-- | Function | Algorithm | Reference | +-- +===================+==================================+===============================+ +-- | 'lu' | LU with partial pivoting | GVL4 Algorithm 3.4.1, p. 126 | +-- +-------------------+----------------------------------+-------------------------------+ +-- | 'luNoPivot' | Outer-product LU (no pivoting) | GVL4 Algorithm 3.2.1, p. 115 | +-- +-------------------+----------------------------------+-------------------------------+ +-- | 'luSolve' | Solve via \(PA = LU\) | GVL4 Section 3.2, p. 118 | +-- +-------------------+----------------------------------+-------------------------------+ +-- | 'det' | Determinant via \(PA = LU\) | GVL4 Section 3.2, p. 120 | +-- +-------------------+----------------------------------+-------------------------------+ +-- +-- == Complexity +-- +-- The factorization requires \(O(2n^3/3)\) flops (GVL4 p. 118). Each +-- subsequent triangular solve adds \(O(n^2)\) flops. +-- +-- == Type Safety +-- +-- Matrix dimensions are tracked at the type level via 'KnownNat', so the +-- compiler statically enforces that the coefficient matrix is square and +-- that right-hand side vectors have conforming length. +module Numeric.LinearAlgebra.Massiv.Solve.LU + ( -- * LU factorization + lu + , luNoPivot + -- * Solving with LU (\(Ax = b\)) + , luSolve + , luSolveP + -- * Determinant + , det + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..), unwrapByteArray, unwrapByteArrayOffset, + unwrapMutableByteArray, unwrapMutableByteArrayOffset) +import GHC.TypeNats (KnownNat) +import Data.Ord (comparing) +import Data.List (maximumBy) +import Control.Monad (when, forM) +import GHC.Exts +import GHC.Prim +import GHC.ST (ST(..)) +import Data.Primitive.ByteArray (ByteArray(..), MutableByteArray(..), newByteArray, unsafeFreezeByteArray) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Solve.Triangular (forwardSubUnit, backSub) +import Numeric.LinearAlgebra.Massiv.Internal.Kernel + (rawLUEliminateColumn, rawLUEliminateColumnTo, rawSwapRows, rawPivotSearch, + rawForwardSubUnitPackedSIMD, rawBackSubPackedSIMD, + rawGemmKernel, rawZeroDoubles) + +-- | LU factorization with partial pivoting (GVL4 Algorithm 3.4.1, p. 126). +-- +-- Given an \(n \times n\) matrix \(A\), computes the factorization +-- +-- \[ +-- PA = LU +-- \] +-- +-- where +-- +-- * \(P\) is a permutation matrix (returned as a pivot-index vector of type +-- @Array P Ix1 Int@), +-- * \(L\) is unit lower triangular (stored /below/ the diagonal of the +-- returned packed matrix), and +-- * \(U\) is upper triangular (stored /on and above/ the diagonal). +-- +-- Partial pivoting selects the entry of largest absolute value in the +-- current column as the pivot, guaranteeing existence for any nonsingular +-- \(A\) (GVL4 Theorem 3.4.1, p. 125). +-- +-- ==== Type-safety guarantees +-- +-- The 'KnownNat' constraint on \(n\) statically ensures the matrix is +-- square. The 'Ord' constraint is required for pivot selection. +-- +-- ==== Complexity +-- +-- \(O(2n^3/3)\) flops (GVL4 p. 118). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 3.4.1 +-- (Outer Product LU with Partial Pivoting), p. 126. +lu :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e, Ord e) + => Matrix n n r e -> (Matrix n n r e, M.Array M.P Ix1 Int) +lu (MkMatrix a) = + let nn = dimVal @n + (pivArr, luArr) = M.withMArrayST a $ \ma -> do + piv <- M.newMArray @M.P (M.Sz1 nn) 0 + mapM_ (\i -> M.write_ piv i i) [0..nn-1] + + mapM_ (\k -> do + -- Find pivot row + vals <- mapM (\i -> do + v <- M.readM ma (i :. k) + pure (i, abs v) + ) [k..nn-1] + let (pivRow, _) = maximumBy (comparing snd) vals + + -- Swap rows k and pivRow + condM (pivRow /= k) $ do + pk <- M.readM piv k + pp <- M.readM piv pivRow + M.write_ piv k pp + M.write_ piv pivRow pk + mapM_ (\j -> do + akj <- M.readM ma (k :. j) + apj <- M.readM ma (pivRow :. j) + M.write_ ma (k :. j) apj + M.write_ ma (pivRow :. j) akj + ) [0..nn-1] + + -- Compute multipliers and update submatrix + akk <- M.readM ma (k :. k) + condM (akk /= 0) $ + mapM_ (\i -> do + aik <- M.readM ma (i :. k) + let mult = aik / akk + M.write_ ma (i :. k) mult + mapM_ (\j -> do + aij <- M.readM ma (i :. j) + akj <- M.readM ma (k :. j) + M.write_ ma (i :. j) (aij - mult * akj) + ) [k+1..nn-1] + ) [k+1..nn-1] + ) [0..nn-2] + + M.freezeS piv + + in (MkMatrix luArr, pivArr) + +-- | LU factorization without pivoting (GVL4 Algorithm 3.2.1, p. 115). +-- +-- Given an \(n \times n\) matrix \(A\) whose leading principal submatrices +-- are all nonsingular, computes the factorization \(A = LU\) where \(L\) is +-- unit lower triangular and \(U\) is upper triangular. Both factors are +-- packed into a single returned matrix: \(L\) occupies the strictly lower +-- triangle (the unit diagonal is implicit) and \(U\) occupies the upper +-- triangle including the diagonal. +-- +-- __Precondition.__ All leading principal submatrices +-- \(A(1{:}k, 1{:}k)\), \(k = 1, \ldots, n\), must be nonsingular +-- (GVL4 Theorem 3.2.1, p. 116). Violating this precondition results in +-- division by zero. +-- +-- ==== Type-safety guarantees +-- +-- The 'KnownNat' constraint on \(n\) statically ensures the input is a +-- square matrix. +-- +-- ==== Complexity +-- +-- \(O(2n^3/3)\) flops (GVL4 p. 118). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 3.2.1 +-- (Outer Product LU Factorization), p. 115. +luNoPivot :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e) + => Matrix n n r e -> Matrix n n r e +luNoPivot (MkMatrix a) = + let nn = dimVal @n + in MkMatrix $ snd $ M.withMArrayST a $ \ma -> + mapM_ (\k -> do + akk <- M.readM ma (k :. k) + mapM_ (\i -> do + aik <- M.readM ma (i :. k) + M.write_ ma (i :. k) (aik / akk) + ) [k+1..nn-1] + mapM_ (\j -> + mapM_ (\i -> do + aij <- M.readM ma (i :. j) + aik <- M.readM ma (i :. k) + akj <- M.readM ma (k :. j) + M.write_ ma (i :. j) (aij - aik * akj) + ) [k+1..nn-1] + ) [k+1..nn-1] + ) [0..nn-2] + +-- | Solve \(Ax = b\) using LU factorization with partial pivoting +-- (GVL4 Section 3.2, p. 118). +-- +-- The algorithm proceeds in three stages: +-- +-- 1. Factor \(PA = LU\) via 'lu' (Algorithm 3.4.1). +-- 2. Solve \(Ly = Pb\) by forward substitution ('forwardSubUnit'). +-- 3. Solve \(Ux = y\) by back substitution ('backSub'). +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) enforces that \(A\) is \(n \times n\) and \(b\) has +-- length \(n\) at compile time. +-- +-- ==== Complexity +-- +-- \(O(2n^3/3)\) flops for the factorization plus \(O(n^2)\) flops for each +-- of the two triangular solves (GVL4 p. 118). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Section 3.2, +-- pp. 114--120. +luSolve :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e, Ord e) + => Matrix n n r e -> Vector n r e -> Vector n r e +luSolve a b = + let (luMat, pivArr) = lu a + -- Extract L (unit lower triangular) + l = makeMatrix @n @n @r $ \i j -> + if i == j then 1 + else if i > j then luMat ! (i, j) + else 0 + -- Extract U (upper triangular) + u = makeMatrix @n @n @r $ \i j -> + if i <= j then luMat ! (i, j) + else 0 + -- Apply permutation to b: pb = P·b + pb = makeVector @n @r $ \i -> + b !. M.index' pivArr i + -- Solve Ly = Pb + y = forwardSubUnit l pb + -- Solve Ux = y + in backSub u y + +-- | Specialised LU solve for @P Double@. +-- Does LU factorisation + solve entirely using raw ByteArray# primops, +-- avoiding L/U matrix reconstruction and massiv's per-element overhead. +-- For n >= 64, uses panel (blocked) LU factorisation with GEMM trailing update. +luSolveP :: forall n. KnownNat n + => Matrix n n M.P Double -> Vector n M.P Double -> Vector n M.P Double +luSolveP (MkMatrix a) (MkVector b) = + let nn = dimVal @n + in createVector @n @M.P $ \mx -> do + -- Thaw matrix for in-place LU factorisation + ma <- M.thawS a + let mbaA = unwrapMutableByteArray ma + offA = unwrapMutableByteArrayOffset ma + + -- Phase 1: LU factorisation with partial pivoting + pivots <- if nn >= 64 + then panelLUFactor mbaA offA nn 32 + else mapM (\k -> do + pivRow <- rawPivotSearch mbaA offA nn k k + condM (pivRow /= k) $ + rawSwapRows mbaA offA nn k pivRow 0 + rawLUEliminateColumn mbaA offA nn k + pure (k, pivRow) + ) [0..nn-2] + + -- Phase 2: Freeze LU and prepare RHS + frozenLU <- M.freezeS ma + let baLU = unwrapByteArray frozenLU + offLU = unwrapByteArrayOffset frozenLU + + -- Copy b into the output vector mx + let mbaX = unwrapMutableByteArray mx + offX = unwrapMutableByteArrayOffset mx + copyVector b mbaX offX nn + + -- Apply pivot permutation to x + applyPivotsForward mbaX offX pivots + + -- Phase 3: Forward substitution (unit lower triangular, SIMD dot-product) + rawForwardSubUnitPackedSIMD baLU offLU nn mbaX offX + + -- Phase 4: Back substitution (upper triangular, SIMD dot-product) + rawBackSubPackedSIMD baLU offLU nn mbaX offX +{-# NOINLINE luSolveP #-} + +-- | Panel (blocked) LU factorisation with GEMM trailing update. +-- Processes columns in panels of width @nb@. Within each panel, elimination +-- is restricted to the panel columns; the trailing submatrix is updated via a +-- single GEMM call, converting O(n) column-by-column Level-2 updates into one +-- cache-friendly Level-3 GEMM. +panelLUFactor :: MutableByteArray s -> Int -> Int -> Int -> ST s [(Int, Int)] +panelLUFactor mbaA offA nn nb = go 0 [] + where + go !k0 !pivAcc + | k0 >= nn - 1 = pure (reverse pivAcc) + | otherwise = do + let !panelEnd = min (k0 + nb) nn + !actualNb = panelEnd - k0 + + -- Factor panel columns k0..panelEnd-1 with restricted trailing update + panelPivs <- forM [k0..panelEnd-1] $ \k -> do + pivRow <- rawPivotSearch mbaA offA nn k k + condM (pivRow /= k) $ + rawSwapRows mbaA offA nn k pivRow 0 + rawLUEliminateColumnTo mbaA offA nn k panelEnd + pure (k, pivRow) + + -- Apply panel's L to trailing columns: triangular solve for U12 + when (panelEnd < nn) $ do + rawTriSolvePanelTrail mbaA offA nn k0 panelEnd + + -- GEMM update: A22 -= L21 × U12 + let !mTrail = nn - panelEnd + !nTrail = nn - panelEnd + + -- Copy L21 to dense buffer (mTrail × actualNb) + bufL <- newByteArray (mTrail * actualNb * 8) + rawCopySubmatrixToDense mbaA offA nn panelEnd k0 mTrail actualNb bufL 0 + + -- Copy U12 to dense buffer (actualNb × nTrail) + bufU <- newByteArray (actualNb * nTrail * 8) + rawCopySubmatrixToDense mbaA offA nn k0 panelEnd actualNb nTrail bufU 0 + + -- Freeze for immutable GEMM inputs + baL <- unsafeFreezeByteArray bufL + baU <- unsafeFreezeByteArray bufU + + -- GEMM: C = L21 × U12 + bufC <- newByteArray (mTrail * nTrail * 8) + rawZeroDoubles bufC 0 (mTrail * nTrail) + rawGemmKernel baL 0 baU 0 bufC 0 mTrail actualNb nTrail + + -- Subtract C from A22 + baC <- unsafeFreezeByteArray bufC + rawSubtractFromStrided baC 0 nTrail mbaA offA nn panelEnd panelEnd mTrail nTrail + + go panelEnd (reverse panelPivs ++ pivAcc) + +-- | Apply unit lower triangular solve from the panel to trailing columns. +-- After factoring panel columns [k0..panelEnd-1] with restricted updates, +-- the trailing columns [panelEnd..n-1] need: for each k in the panel, apply +-- the multipliers to rows k+1..panelEnd-1 of the trailing columns. +-- i.e. A[i,j] -= A[i,k] * A[k,j] for k0 <= k < panelEnd, k < i < panelEnd, panelEnd <= j < n +rawTriSolvePanelTrail :: MutableByteArray s -> Int -> Int -> Int -> Int -> ST s () +rawTriSolvePanelTrail (MutableByteArray mba) (I# off) (I# n) (I# k0) (I# panelEnd) = ST $ \s0 -> + let -- For each column k in the panel + goK k s + | isTrue# (k >=# panelEnd) = s + | otherwise = + let kRowOff = off +# k *# n + in goI k (k +# 1#) kRowOff s + where + -- For each row i in [k+1..panelEnd-1] + goI k_ i kRowOff s_ + | isTrue# (i >=# panelEnd) = goK (k_ +# 1#) s_ + | otherwise = + let iRowOff = off +# i *# n + in case readDoubleArray# mba (iRowOff +# k_) s_ of + (# s1, lik #) -> + let negLik = negateDouble# lik + negLikV = broadcastDoubleX4# negLik + jSpan = n -# panelEnd + j4End = panelEnd +# (jSpan -# (jSpan `remInt#` 4#)) + -- SIMD j-loop + goJSimd j s__ + | isTrue# (j >=# j4End) = s__ + | otherwise = + case readDoubleArrayAsDoubleX4# mba (iRowOff +# j) s__ of + (# s2, aij #) -> + case readDoubleArrayAsDoubleX4# mba (kRowOff +# j) s2 of + (# s3, akj #) -> + let aij' = fmaddDoubleX4# negLikV akj aij + in case writeDoubleArrayAsDoubleX4# mba (iRowOff +# j) aij' s3 of + s4 -> goJSimd (j +# 4#) s4 + -- Scalar cleanup + goJScalar j s__ + | isTrue# (j >=# n) = s__ + | otherwise = + case readDoubleArray# mba (iRowOff +# j) s__ of + (# s2, aij #) -> + case readDoubleArray# mba (kRowOff +# j) s2 of + (# s3, akj #) -> + case writeDoubleArray# mba (iRowOff +# j) (aij +## negLik *## akj) s3 of + s4 -> goJScalar (j +# 1#) s4 + in goI k_ (i +# 1#) kRowOff (goJScalar j4End (goJSimd panelEnd s1)) + in (# goK k0 s0, () #) +{-# INLINE rawTriSolvePanelTrail #-} + +-- | Copy a submatrix A[rowStart..rowStart+m-1, colStart..colStart+k-1] (stride n) +-- into a dense buffer (stride k). +rawCopySubmatrixToDense :: MutableByteArray s -> Int -> Int -- src, offset, n + -> Int -> Int -> Int -> Int -- rowStart, colStart, m, k + -> MutableByteArray s -> Int -- dst, dstOffset + -> ST s () +rawCopySubmatrixToDense (MutableByteArray mba_src) (I# off_src) (I# n) + (I# rowStart) (I# colStart) (I# m) (I# k) + (MutableByteArray mba_dst) (I# off_dst) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# m) = s + | otherwise = + let srcRow = off_src +# (rowStart +# i) *# n +# colStart + dstRow = off_dst +# i *# k + span_ = k -# (k `remInt#` 4#) + goSimd j s_ + | isTrue# (j >=# span_) = s_ + | otherwise = + case readDoubleArrayAsDoubleX4# mba_src (srcRow +# j) s_ of + (# s1, v #) -> + case writeDoubleArrayAsDoubleX4# mba_dst (dstRow +# j) v s1 of + s2 -> goSimd (j +# 4#) s2 + goScalar j s_ + | isTrue# (j >=# k) = s_ + | otherwise = + case readDoubleArray# mba_src (srcRow +# j) s_ of + (# s1, v #) -> + case writeDoubleArray# mba_dst (dstRow +# j) v s1 of + s2 -> goScalar (j +# 1#) s2 + in goI (i +# 1#) (goScalar span_ (goSimd 0# s)) + in (# goI 0# s0, () #) +{-# INLINE rawCopySubmatrixToDense #-} + +-- | Subtract a dense matrix C (m × k, stride srcStride) from a strided submatrix +-- A[rowStart..rowStart+m-1, colStart..colStart+k-1] (stride n). +rawSubtractFromStrided :: ByteArray -> Int -> Int -- src, srcOffset, srcStride + -> MutableByteArray s -> Int -> Int -- dst, dstOffset, n + -> Int -> Int -> Int -> Int -- rowStart, colStart, m, k + -> ST s () +rawSubtractFromStrided (ByteArray ba_src) (I# off_src) (I# srcStride) + (MutableByteArray mba_dst) (I# off_dst) (I# n) + (I# rowStart) (I# colStart) (I# m) (I# k) = ST $ \s0 -> + let goI i s + | isTrue# (i >=# m) = s + | otherwise = + let srcRow = off_src +# i *# srcStride + dstRow = off_dst +# (rowStart +# i) *# n +# colStart + span_ = k -# (k `remInt#` 4#) + goSimd j s_ + | isTrue# (j >=# span_) = s_ + | otherwise = + case readDoubleArrayAsDoubleX4# mba_dst (dstRow +# j) s_ of + (# s1, aij #) -> + let cij = indexDoubleArrayAsDoubleX4# ba_src (srcRow +# j) + aij' = plusDoubleX4# aij (negateDoubleX4# cij) + in case writeDoubleArrayAsDoubleX4# mba_dst (dstRow +# j) aij' s1 of + s2 -> goSimd (j +# 4#) s2 + goScalar j s_ + | isTrue# (j >=# k) = s_ + | otherwise = + case readDoubleArray# mba_dst (dstRow +# j) s_ of + (# s1, aij #) -> + let cij = indexDoubleArray# ba_src (srcRow +# j) + in case writeDoubleArray# mba_dst (dstRow +# j) (aij -## cij) s1 of + s2 -> goScalar (j +# 1#) s2 + in goI (i +# 1#) (goScalar span_ (goSimd 0# s)) + in (# goI 0# s0, () #) +{-# INLINE rawSubtractFromStrided #-} + +-- | Copy an immutable P vector into a mutable byte array. +copyVector :: M.Array M.P Ix1 Double -> MutableByteArray s -> Int -> Int -> ST s () +copyVector src (MutableByteArray mba_dst) (I# off_dst) (I# n) = ST $ \s0 -> + let ba_src = case unwrapByteArray src of ByteArray ba -> ba + off_src = case unwrapByteArrayOffset src of I# o -> o + go i s + | isTrue# (i >=# n) = s + | otherwise = + let v = indexDoubleArray# ba_src (off_src +# i) + in case writeDoubleArray# mba_dst (off_dst +# i) v s of + s' -> go (i +# 1#) s' + in (# go 0# s0, () #) +{-# INLINE copyVector #-} + +-- | Apply pivot permutation to a vector (forward direction). +applyPivotsForward :: MutableByteArray s -> Int -> [(Int, Int)] -> ST s () +applyPivotsForward (MutableByteArray mba) (I# off) pivots = ST $ \s0 -> + let go [] s = s + go ((I# k, I# pivRow) : rest) s + | isTrue# (k ==# pivRow) = go rest s + | otherwise = + case readDoubleArray# mba (off +# k) s of + (# s1, vk #) -> + case readDoubleArray# mba (off +# pivRow) s1 of + (# s2, vp #) -> + case writeDoubleArray# mba (off +# k) vp s2 of + s3 -> case writeDoubleArray# mba (off +# pivRow) vk s3 of + s4 -> go rest s4 + in (# go pivots s0, () #) +{-# INLINE applyPivotsForward #-} + +-- | Compute the determinant of an \(n \times n\) matrix via LU factorization +-- (GVL4 Section 3.2, p. 120). +-- +-- ==== Mathematical definition +-- +-- From \(PA = LU\) it follows that +-- +-- \[ +-- \det(A) = (-1)^s \prod_{i=1}^{n} u_{ii} +-- \] +-- +-- where \(s\) is the number of row transpositions performed during partial +-- pivoting. +-- +-- ==== Type-safety guarantees +-- +-- 'KnownNat' \(n\) ensures \(A\) is square at compile time. +-- +-- ==== Complexity +-- +-- \(O(2n^3/3)\) flops, dominated by the LU factorization. +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Section 3.2, p. 120. +det :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e, Ord e) + => Matrix n n r e -> e +det a = + let nn = dimVal @n + (luMat, pivArr) = lu a + -- Product of U diagonal + diagProd = foldl' (\acc i -> acc * (luMat ! (i, i))) 1 [0..nn-1] + -- Count transpositions: number of i where piv[i] /= i + pivList = map (M.index' pivArr) [0..nn-1] + nswaps = countSwaps pivList + sign = if even nswaps then 1 else -1 + in sign * diagProd + +-- | Count the number of swaps in a permutation. +countSwaps :: [Int] -> Int +countSwaps perm = go (zip [0..] perm) 0 + where + go [] n = n + go ((i,p):rest) n + | i == p = go rest n + | otherwise = + -- Swap p into position i by finding where i is + let rest' = map (\(idx, v) -> if v == i then (idx, p) else (idx, v)) rest + in go rest' (n + 1) + +-- | Conditional monadic action. +condM :: Applicative m => Bool -> m () -> m () +condM True act = act +condM False _ = pure () diff --git a/src/Numeric/LinearAlgebra/Massiv/Solve/Triangular.hs b/src/Numeric/LinearAlgebra/Massiv/Solve/Triangular.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Solve/Triangular.hs @@ -0,0 +1,250 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Solve.Triangular +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- = Triangular System Solvers +-- +-- Forward and back substitution for lower- and upper-triangular linear +-- systems, following Golub & Van Loan, /Matrix Computations/, 4th edition +-- (GVL4), Section 3.1, pp. 106--113. +-- +-- The section presents four algorithmic variants organised along two axes: +-- +-- * __Row-oriented vs. column-oriented__ inner loops, affecting data-access +-- patterns and cache behaviour. +-- * __General vs. unit-triangular__ coefficient matrices, where the unit +-- variants avoid division by the (known-to-be-one) diagonal. +-- +-- This module exposes the following mapping: +-- +-- +-------------------+---------------------------+---------------------------------+ +-- | Function | Algorithm | Reference | +-- +===================+===========================+=================================+ +-- | 'forwardSub' | Row-oriented forward sub | GVL4 Algorithm 3.1.1, p. 106 | +-- +-------------------+---------------------------+---------------------------------+ +-- | 'backSub' | Row-oriented back sub | GVL4 Algorithm 3.1.2, p. 107 | +-- +-------------------+---------------------------+---------------------------------+ +-- | 'forwardSubUnit' | Column-oriented fwd sub | GVL4 Algorithm 3.1.3, p. 108 | +-- +-------------------+---------------------------+---------------------------------+ +-- | 'backSubUnit' | Column-oriented back sub | GVL4 Algorithm 3.1.4, p. 109 | +-- +-------------------+---------------------------+---------------------------------+ +-- +-- == Complexity +-- +-- Every solver performs \(O(n^2 / 2)\) floating-point operations (flops) for +-- an \(n \times n\) triangular system (GVL4 p. 109). +-- +-- == Type Safety +-- +-- The matrix dimension \(n\) is tracked at the type level via 'KnownNat', +-- so the compiler statically guarantees that the coefficient matrix is +-- square and that the right-hand side vector has a conforming length. +module Numeric.LinearAlgebra.Massiv.Solve.Triangular + ( -- * Forward substitution (\(Lx = b\)) + forwardSub + -- * Back substitution (\(Ux = b\)) + , backSub + -- * Unit-triangular variants + , forwardSubUnit + , backSubUnit + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Sz(..)) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal + +-- | Row-oriented forward substitution (GVL4 Algorithm 3.1.1, p. 106). +-- +-- Solves the lower-triangular system \(Lx = b\) where \(L \in \mathbb{R}^{n \times n}\) +-- is lower triangular with nonzero diagonal entries. +-- +-- ==== Mathematical definition +-- +-- For \(j = 1, \ldots, n\): +-- +-- \[ +-- x_j = \frac{1}{\ell_{jj}} \left( b_j - \sum_{k=1}^{j-1} \ell_{jk}\, x_k \right) +-- \] +-- +-- ==== Type-safety guarantees +-- +-- The type-level natural \(n\) ('KnownNat') ensures that \(L\) is square +-- and that \(b\) has exactly \(n\) entries. A dimension mismatch is a +-- compile-time error. +-- +-- ==== Complexity +-- +-- \(O(n^2 / 2)\) flops (GVL4 p. 109). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 3.1.1 +-- (Row-Oriented Forward Substitution), p. 106. +forwardSub :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e) + => Matrix n n r e -> Vector n r e -> Vector n r e +forwardSub l b = createVector @n $ \mx -> do + let nn = dimVal @n + -- Copy b into the mutable result + mapM_ (\i -> M.write_ mx i (b !. i)) [0..nn-1] + -- Forward elimination + mapM_ (\j -> do + xj <- M.readM mx j + let ldiag = l ! (j, j) + xj' = xj / ldiag + M.write_ mx j xj' + -- Update remaining entries + mapM_ (\i -> do + xi <- M.readM mx i + M.write_ mx i (xi - (l ! (i, j)) * xj') + ) [j+1..nn-1] + ) [0..nn-1] + +-- | Row-oriented back substitution (GVL4 Algorithm 3.1.2, p. 107). +-- +-- Solves the upper-triangular system \(Ux = b\) where \(U \in \mathbb{R}^{n \times n}\) +-- is upper triangular with nonzero diagonal entries. +-- +-- ==== Mathematical definition +-- +-- For \(j = n, n-1, \ldots, 1\): +-- +-- \[ +-- x_j = \frac{1}{u_{jj}} \left( b_j - \sum_{k=j+1}^{n} u_{jk}\, x_k \right) +-- \] +-- +-- ==== Type-safety guarantees +-- +-- The type-level natural \(n\) ('KnownNat') ensures that \(U\) is square +-- and that \(b\) has exactly \(n\) entries. A dimension mismatch is a +-- compile-time error. +-- +-- ==== Complexity +-- +-- \(O(n^2 / 2)\) flops (GVL4 p. 109). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 3.1.2 +-- (Row-Oriented Back Substitution), p. 107. +backSub :: forall n r e. (KnownNat n, M.Manifest r e, Fractional e) + => Matrix n n r e -> Vector n r e -> Vector n r e +backSub u b = createVector @n $ \mx -> do + let nn = dimVal @n + -- Copy b into the mutable result + mapM_ (\i -> M.write_ mx i (b !. i)) [0..nn-1] + -- Back elimination + mapM_ (\j -> do + xj <- M.readM mx j + let udiag = u ! (j, j) + xj' = xj / udiag + M.write_ mx j xj' + -- Update remaining entries + mapM_ (\i -> do + xi <- M.readM mx i + M.write_ mx i (xi - (u ! (i, j)) * xj') + ) [0..j-1] + ) [nn-1, nn-2..0] + +-- | Column-oriented forward substitution for unit lower triangular systems +-- (GVL4 Algorithm 3.1.3, p. 108). +-- +-- Solves \(Lx = b\) where \(L \in \mathbb{R}^{n \times n}\) is /unit/ lower +-- triangular, i.e. \(\ell_{jj} = 1\) for all \(j\). Because the diagonal is +-- implicitly one, no division is needed and the constraint relaxes from +-- 'Fractional' to 'Num'. +-- +-- ==== Mathematical definition +-- +-- For \(j = 1, \ldots, n\): +-- +-- \[ +-- x_j = b_j - \sum_{k=1}^{j-1} \ell_{jk}\, x_k +-- \] +-- +-- The implementation uses a /column-oriented/ (saxpy) loop: once \(x_j\) is +-- determined, rows \(i > j\) are updated by subtracting \(\ell_{ij}\, x_j\). +-- +-- ==== Type-safety guarantees +-- +-- Identical to 'forwardSub': the dimensions are enforced at compile time +-- via 'KnownNat'. +-- +-- ==== Complexity +-- +-- \(O(n^2 / 2)\) flops (GVL4 p. 109). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 3.1.3 +-- (Column-Oriented Forward Substitution), p. 108. +forwardSubUnit :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Matrix n n r e -> Vector n r e -> Vector n r e +forwardSubUnit l b = createVector @n $ \mx -> do + let nn = dimVal @n + -- Copy b into the mutable result + mapM_ (\i -> M.write_ mx i (b !. i)) [0..nn-1] + -- Column-oriented forward substitution + mapM_ (\j -> do + xj <- M.readM mx j + -- Subtract l(i,j)*xj from x(i) for i > j + mapM_ (\i -> do + xi <- M.readM mx i + M.write_ mx i (xi - (l ! (i, j)) * xj) + ) [j+1..nn-1] + ) [0..nn-1] + +-- | Column-oriented back substitution for unit upper triangular systems +-- (GVL4 Algorithm 3.1.4, p. 109). +-- +-- Solves \(Ux = b\) where \(U \in \mathbb{R}^{n \times n}\) is /unit/ upper +-- triangular, i.e. \(u_{jj} = 1\) for all \(j\). Because the diagonal is +-- implicitly one, no division is needed and the constraint relaxes from +-- 'Fractional' to 'Num'. +-- +-- ==== Mathematical definition +-- +-- For \(j = n, n-1, \ldots, 1\): +-- +-- \[ +-- x_j = b_j - \sum_{k=j+1}^{n} u_{jk}\, x_k +-- \] +-- +-- The implementation uses a /column-oriented/ loop: once \(x_j\) is known, +-- rows \(i < j\) are updated by subtracting \(u_{ij}\, x_j\). +-- +-- ==== Type-safety guarantees +-- +-- Identical to 'backSub': the dimensions are enforced at compile time via +-- 'KnownNat'. +-- +-- ==== Complexity +-- +-- \(O(n^2 / 2)\) flops (GVL4 p. 109). +-- +-- ==== Reference +-- +-- Golub & Van Loan, /Matrix Computations/, 4th ed., Algorithm 3.1.4 +-- (Column-Oriented Back Substitution), p. 109. +backSubUnit :: forall n r e. (KnownNat n, M.Manifest r e, Num e) + => Matrix n n r e -> Vector n r e -> Vector n r e +backSubUnit u b = createVector @n $ \mx -> do + let nn = dimVal @n + -- Copy b into the mutable result + mapM_ (\i -> M.write_ mx i (b !. i)) [0..nn-1] + -- Column-oriented back substitution + mapM_ (\j -> do + xj <- M.readM mx j + -- Subtract u(i,j)*xj from x(i) for i < j + mapM_ (\i -> do + xi <- M.readM mx i + M.write_ mx i (xi - (u ! (i, j)) * xj) + ) [0..j-1] + ) [nn-1, nn-2..0] diff --git a/src/Numeric/LinearAlgebra/Massiv/Types.hs b/src/Numeric/LinearAlgebra/Massiv/Types.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/LinearAlgebra/Massiv/Types.hs @@ -0,0 +1,185 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | +-- Module : Numeric.LinearAlgebra.Massiv.Types +-- Copyright : (c) Nadia Chambers 2026 +-- License : BSD-3-Clause +-- Maintainer : nadia.chambers@iohk.io +-- Stability : experimental +-- +-- Core type definitions for type-safe dimensioned matrices and vectors +-- backed by arrays. +-- +-- The central types are 'Matrix' and 'Vector', which wrap massiv arrays +-- with phantom type-level natural number parameters encoding their dimensions. +-- This ensures that dimensionally incorrect operations (e.g., adding matrices +-- of different sizes, or multiplying matrices with incompatible inner dimensions) +-- are caught at compile time by GHC's type checker. +-- +-- = Type-level dimension encoding +-- +-- Dimensions are encoded as GHC @DataKinds@ promoted @'GHC.TypeNats.Nat'@ +-- values. The constraint @'GHC.TypeNats.KnownNat' n@ provides access to the +-- runtime value via @'GHC.TypeNats.natVal'@. +-- +-- @ +-- -- A 3x4 matrix of Doubles using Primitive representation +-- type MyMatrix = Matrix 3 4 P Double +-- +-- -- A 5-element vector of Doubles using Unboxed representation +-- type MyVector = Vector 5 U Double +-- @ +-- +-- = Existential wrappers +-- +-- For situations where dimensions are not known until runtime (e.g., reading +-- a matrix from a file), use 'SomeMatrix' and 'SomeVector'. These existentially +-- quantify the dimension parameters while retaining 'KnownNat' evidence. +-- +-- See "Numeric.LinearAlgebra.Massiv.Internal" for construction helpers. +module Numeric.LinearAlgebra.Massiv.Types + ( -- * Dimensioned matrix type + Matrix(..) + -- * Dimensioned vector type + , Vector(..) + -- * Smart constructors + , matrix + , vector + -- * Existential wrappers + , SomeMatrix(..) + , SomeVector(..) + , someMatrix + , someVector + -- * Dimension queries + , rows + , cols + , size + -- * Type-level helpers + , type KnownDims + ) where + +import Data.Massiv.Array (Array, Ix2(..), Sz(..), Ix1, Comp(..)) +import qualified Data.Massiv.Array as M +import GHC.TypeNats (Nat, KnownNat, natVal, SomeNat(..), someNatVal) +import Data.Proxy (Proxy(..)) +import Control.DeepSeq (NFData(..)) + +-- | Constraint synonym for two known dimensions. +-- +-- @KnownDims m n@ is equivalent to @(KnownNat m, KnownNat n)@. +type KnownDims m n = (KnownNat m, KnownNat n) + +-- | A matrix with compile-time known dimensions \(m\) (rows) \(\times\) \(n\) (cols). +-- +-- Wraps a massiv @'Data.Massiv.Array.Array' r 'Data.Massiv.Array.Ix2' e@. +-- The phantom type parameters \(m\) and \(n\) enforce dimensional conformance +-- at compile time. For example, matrix multiplication via 'matMul' requires +-- the inner dimensions to unify: +-- +-- @ +-- 'matMul' :: Matrix m __k__ r e -> Matrix __k__ n r e -> Matrix m n r e +-- @ +-- +-- The representation parameter @r@ selects the massiv array backend: +-- +-- * @'Data.Massiv.Array.P'@ — Primitive (best for 'Double', 'Int'; pinned memory) +-- * @'Data.Massiv.Array.U'@ — Unboxed (via @Data.Vector.Unboxed@) +-- * @'Data.Massiv.Array.S'@ — Storable (via @Foreign.ForeignPtr@; useful for FFI) +-- * @'Data.Massiv.Array.B'@ — Boxed (polymorphic but slower; GC overhead) +newtype Matrix (m :: Nat) (n :: Nat) r e = MkMatrix { unMatrix :: Array r Ix2 e } + +deriving instance Show (Array r Ix2 e) => Show (Matrix m n r e) +deriving instance Eq (Array r Ix2 e) => Eq (Matrix m n r e) + +instance NFData (Array r Ix2 e) => NFData (Matrix m n r e) where + rnf (MkMatrix arr) = rnf arr + +-- | A vector with compile-time known dimension \(n\). +-- +-- Wraps a massiv @'Data.Massiv.Array.Array' r 'Data.Massiv.Array.Ix1' e@. +-- The phantom parameter \(n\) ensures that vector operations (e.g., 'dot', +-- 'axpy') are only applied to vectors of matching dimension. +newtype Vector (n :: Nat) r e = MkVector { unVector :: Array r Ix1 e } + +deriving instance Show (Array r Ix1 e) => Show (Vector n r e) +deriving instance Eq (Array r Ix1 e) => Eq (Vector n r e) + +instance NFData (Array r Ix1 e) => NFData (Vector n r e) where + rnf (MkVector arr) = rnf arr + +-- | Smart constructor for matrices. Checks at runtime that the array +-- dimensions match the type-level dimensions \(m\) and \(n\). +-- +-- Returns 'Nothing' if the dimensions do not match. +-- +-- @ +-- let arr = M.makeArray Seq (Sz2 3 4) (\\(i :. j) -> fromIntegral (i + j)) +-- matrix \@3 \@4 arr -- Just (MkMatrix arr) +-- matrix \@2 \@4 arr -- Nothing +-- @ +matrix :: forall m n r e. (KnownDims m n, M.Size r) + => Array r Ix2 e -> Maybe (Matrix m n r e) +matrix arr + | M.Sz2 r c <- M.size arr + , r == fromIntegral (natVal (Proxy @m)) + , c == fromIntegral (natVal (Proxy @n)) + = Just (MkMatrix arr) + | otherwise = Nothing + +-- | Smart constructor for vectors. Checks at runtime that the array +-- size matches the type-level dimension \(n\). +-- +-- Returns 'Nothing' if the size does not match. +vector :: forall n r e. (KnownNat n, M.Size r) + => Array r Ix1 e -> Maybe (Vector n r e) +vector arr + | M.Sz1 n <- M.size arr + , n == fromIntegral (natVal (Proxy @n)) + = Just (MkVector arr) + | otherwise = Nothing + +-- | Get the number of rows at the value level. \(O(1)\). +rows :: forall m n r e. KnownNat m => Matrix m n r e -> Int +rows _ = fromIntegral (natVal (Proxy @m)) + +-- | Get the number of columns at the value level. \(O(1)\). +cols :: forall m n r e. KnownNat n => Matrix m n r e -> Int +cols _ = fromIntegral (natVal (Proxy @n)) + +-- | Get the size of a vector at the value level. \(O(1)\). +size :: forall n r e. KnownNat n => Vector n r e -> Int +size _ = fromIntegral (natVal (Proxy @n)) + +-- | Existentially quantified matrix with runtime-determined dimensions. +-- +-- Use 'someMatrix' to wrap a massiv array whose dimensions are not known +-- at compile time. Pattern matching on 'SomeMatrix' brings 'KnownNat' +-- evidence into scope: +-- +-- @ +-- case someMatrix arr of +-- SomeMatrix (mat :: Matrix m n r e) -> ... +-- -- m and n are now in scope as KnownNat +-- @ +data SomeMatrix r e where + SomeMatrix :: (KnownNat m, KnownNat n) => Matrix m n r e -> SomeMatrix r e + +-- | Existentially quantified vector with runtime-determined dimensions. +data SomeVector r e where + SomeVector :: KnownNat n => Vector n r e -> SomeVector r e + +-- | Wrap a massiv 2D array into an existentially typed matrix. +someMatrix :: M.Size r => Array r Ix2 e -> SomeMatrix r e +someMatrix arr = + let M.Sz2 r c = M.size arr + in case someNatVal (fromIntegral r) of + SomeNat (_ :: Proxy m) -> case someNatVal (fromIntegral c) of + SomeNat (_ :: Proxy n) -> SomeMatrix @m @n (MkMatrix arr) + +-- | Wrap a massiv 1D array into an existentially typed vector. +someVector :: M.Size r => Array r Ix1 e -> SomeVector r e +someVector arr = + let M.Sz1 n = M.size arr + in case someNatVal (fromIntegral n) of + SomeNat (_ :: Proxy n) -> SomeVector @n (MkVector arr) diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,18 @@ +module Main (main) where + +import Test.Tasty + +import Test.BLAS (blasTests) +import Test.Solve (solveTests) +import Test.Orthogonal (orthogonalTests) +import Test.Eigen (eigenTests) +import Test.Norms (normTests) + +main :: IO () +main = defaultMain $ testGroup "linear-massiv" + [ blasTests + , solveTests + , orthogonalTests + , eigenTests + , normTests + ] diff --git a/test/Test/BLAS.hs b/test/Test/BLAS.hs new file mode 100644 --- /dev/null +++ b/test/Test/BLAS.hs @@ -0,0 +1,172 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Test.BLAS (blasTests) where + +import Test.Tasty +import Test.Tasty.QuickCheck +import Test.Tasty.HUnit +import qualified Data.Massiv.Array as M + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level1 +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 +import Numeric.LinearAlgebra.Massiv.Norms (normFrob) +import Test.Types (genMatrix, genVector, hilbertMatrix, (~=), matApproxEq, vecApproxEq) +import Test.Residuals (machineEps) + +blasTests :: TestTree +blasTests = testGroup "BLAS" + [ testGroup "Level 1" + [ testProperty "dot product commutative" prop_dotCommutative + , testProperty "dot with zero vector" prop_dotZero + , testCase "dot Hilbert columns" test_dotHilbertColumns + , testProperty "axpy identity" prop_axpyIdentity + , testProperty "scal by 1" prop_scalIdentity + , testProperty "nrm2 non-negative" prop_nrm2NonNeg + ] + , testGroup "Level 2" + [ testProperty "matvec with identity" prop_matvecIdentity + , testProperty "gemv alpha=1 beta=0" prop_gemvSimple + ] + , testGroup "Level 3" + [ testProperty "matMul with identity (left)" prop_matMulIdentityLeft + , testProperty "matMul with identity (right)" prop_matMulIdentityRight + , testProperty "transpose involution" prop_transposeInvolution + , testProperty "mAdd commutative" prop_mAddCommutative + , testProperty "matMul associative 5x5" prop_matMulAssociative5 + , testProperty "matMul with identity 10x10" prop_matMulIdentity10 + , testCase "3x3 matmul known" test_matMulKnown + , testCase "gemm larger 3x3" test_gemmLarger + ] + ] + +-- Level 1 properties + +prop_dotCommutative :: Property +prop_dotCommutative = forAll ((,) <$> genVector @4 <*> genVector @4) $ \(x, y) -> + dot x y ~= dot y x + +prop_dotZero :: Property +prop_dotZero = forAll (genVector @4) $ \x -> + let z = zeroVector @4 @M.P :: Vector 4 M.P Double + in dot x z ~= 0 + +prop_axpyIdentity :: Property +prop_axpyIdentity = forAll (genVector @4) $ \x -> + let z = zeroVector @4 @M.P :: Vector 4 M.P Double + in vecApproxEq @4 (axpy 1 z x) x + +prop_scalIdentity :: Property +prop_scalIdentity = forAll (genVector @4) $ \x -> + vecApproxEq @4 (scal 1 x) x + +prop_nrm2NonNeg :: Property +prop_nrm2NonNeg = forAll (genVector @4) $ \x -> + nrm2 x >= (0 :: Double) + +-- Level 2 properties + +prop_matvecIdentity :: Property +prop_matvecIdentity = forAll (genVector @3) $ \x -> + let eye = identityMatrix @3 @M.P :: Matrix 3 3 M.P Double + in vecApproxEq @3 (matvec eye x) x + +prop_gemvSimple :: Property +prop_gemvSimple = forAll ((,) <$> genMatrix @3 @3 <*> genVector @3) $ \(a, x) -> + let z = zeroVector @3 @M.P :: Vector 3 M.P Double + result = gemv 1.0 a x 0.0 z + expected = matvec a x + in vecApproxEq @3 result expected + +-- Level 3 properties + +prop_matMulIdentityLeft :: Property +prop_matMulIdentityLeft = forAll (genMatrix @3 @3) $ \a -> + let eye = identityMatrix @3 @M.P :: Matrix 3 3 M.P Double + in matApproxEq @3 @3 (matMul eye a) a + +prop_matMulIdentityRight :: Property +prop_matMulIdentityRight = forAll (genMatrix @3 @3) $ \a -> + let eye = identityMatrix @3 @M.P :: Matrix 3 3 M.P Double + in matApproxEq @3 @3 (matMul a eye) a + +prop_transposeInvolution :: Property +prop_transposeInvolution = forAll (genMatrix @3 @4) $ \a -> + matApproxEq @3 @4 (transpose (transpose a)) a + +prop_mAddCommutative :: Property +prop_mAddCommutative = forAll ((,) <$> genMatrix @3 @3 <*> genMatrix @3 @3) $ \(a, b) -> + matApproxEq @3 @3 (mAdd a b) (mAdd b a) + +-- | (AB)C ≈ A(BC) for random 5×5 matrices. +prop_matMulAssociative5 :: Property +prop_matMulAssociative5 = + forAll ((,,) <$> genMatrix @5 @5 <*> genMatrix @5 @5 <*> genMatrix @5 @5) $ \(a, b, c) -> + let lhs = matMul (matMul a b) c + rhs = matMul a (matMul b c) + in normFrob (mSub lhs rhs) / (normFrob lhs + 1e-15) < 1e-6 + +-- | I·A = A for 10×10 matrices. +prop_matMulIdentity10 :: Property +prop_matMulIdentity10 = forAll (genMatrix @10 @10) $ \a -> + let eye = identityMatrix @10 @M.P :: Matrix 10 10 M.P Double + in matApproxEq @10 @10 (matMul eye a) a + +-- | Dot product of columns 0 and 1 of hilbertMatrix @5. +-- col0 = [1, 1/2, 1/3, 1/4, 1/5] +-- col1 = [1/2, 1/3, 1/4, 1/5, 1/6] +-- dot = sum_{k=0}^{4} 1/((k+1)*(k+2)) = 1/2 + 1/6 + 1/12 + 1/20 + 1/30 = 50/60 = 5/6 +test_dotHilbertColumns :: Assertion +test_dotHilbertColumns = do + let h = hilbertMatrix @5 :: Matrix 5 5 M.P Double + col0 = makeVector @5 @M.P $ \k -> h ! (k, 0) + col1 = makeVector @5 @M.P $ \k -> h ! (k, 1) + result = dot col0 col1 + expected = 5 / 6 :: Double + assertBool ("dot of Hilbert cols 0,1 = 5/6, got " ++ show result) + $ abs (result - expected) < 1e-12 + +-- Known-value test for 3×3 matrix multiplication +test_matMulKnown :: Assertion +test_matMulKnown = do + -- A = [[1,2],[3,4]], B = [[5,6],[7,8]] + -- AB = [[19,22],[43,50]] + let a = makeMatrix @2 @2 @M.P $ \i j -> case (i,j) of + (0,0) -> 1; (0,1) -> 2; (1,0) -> 3; (1,1) -> 4; _ -> 0 :: Double + b = makeMatrix @2 @2 @M.P $ \i j -> case (i,j) of + (0,0) -> 5; (0,1) -> 6; (1,0) -> 7; (1,1) -> 8; _ -> 0 :: Double + c = matMul a b + assertBool "c(0,0) = 19" $ (c ! (0,0)) ~= 19 + assertBool "c(0,1) = 22" $ (c ! (0,1)) ~= 22 + assertBool "c(1,0) = 43" $ (c ! (1,0)) ~= 43 + assertBool "c(1,1) = 50" $ (c ! (1,1)) ~= 50 + +-- | GEMM with α=2.0, β=0.5 on known 3×3 matrices. +-- A = [[1,2,3],[4,5,6],[7,8,9]] +-- B = [[9,8,7],[6,5,4],[3,2,1]] +-- C = [[1,0,0],[0,1,0],[0,0,1]] +-- Result = α*A*B + β*C +-- +-- A*B = [[30,24,18],[84,69,54],[138,114,90]] +-- α*A*B = [[60,48,36],[168,138,108],[276,228,180]] +-- β*C = [[0.5,0,0],[0,0.5,0],[0,0,0.5]] +-- Final = [[60.5,48,36],[168,138.5,108],[276,228,180.5]] +test_gemmLarger :: Assertion +test_gemmLarger = do + let a = makeMatrix @3 @3 @M.P $ \i j -> + fromIntegral (i * 3 + j + 1) :: Double + b = makeMatrix @3 @3 @M.P $ \i j -> + fromIntegral (9 - (i * 3 + j)) :: Double + c = identityMatrix @3 @M.P :: Matrix 3 3 M.P Double + result = gemm 2.0 a b 0.5 c + assertBool "gemm(0,0) = 60.5" $ (result ! (0,0)) ~= 60.5 + assertBool "gemm(0,1) = 48" $ (result ! (0,1)) ~= 48 + assertBool "gemm(0,2) = 36" $ (result ! (0,2)) ~= 36 + assertBool "gemm(1,0) = 168" $ (result ! (1,0)) ~= 168 + assertBool "gemm(1,1) = 138.5"$ (result ! (1,1)) ~= 138.5 + assertBool "gemm(1,2) = 108" $ (result ! (1,2)) ~= 108 + assertBool "gemm(2,0) = 276" $ (result ! (2,0)) ~= 276 + assertBool "gemm(2,1) = 228" $ (result ! (2,1)) ~= 228 + assertBool "gemm(2,2) = 180.5"$ (result ! (2,2)) ~= 180.5 diff --git a/test/Test/Eigen.hs b/test/Test/Eigen.hs new file mode 100644 --- /dev/null +++ b/test/Test/Eigen.hs @@ -0,0 +1,587 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Test.Eigen (eigenTests) where + +import Test.Tasty +import Test.Tasty.QuickCheck +import Test.Tasty.HUnit +import qualified Data.Massiv.Array as M +import Data.List (sort) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvec) +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, matMulP, transpose, mSub) +import Numeric.LinearAlgebra.Massiv.BLAS.Level1 (nrm2, scal) +import Numeric.LinearAlgebra.Massiv.Eigen.Power +import Numeric.LinearAlgebra.Massiv.Eigen.Hessenberg +import Numeric.LinearAlgebra.Massiv.Eigen.Symmetric +import Numeric.LinearAlgebra.Massiv.Eigen.SVD (svd, svdP, svdGKP) +import Data.Proxy (Proxy(..)) +import Numeric.LinearAlgebra.Massiv.Eigen.Schur (schur, eigenvalues) +import Numeric.LinearAlgebra.Massiv.Norms (normFrob, vnorm2) +import Test.Types +import Test.Residuals + +eigenTests :: TestTree +eigenTests = testGroup "Eigenvalue" + [ testGroup "Power Method" + [ testCase "dominant eigenvalue of diagonal" test_powerDiagonal + ] + , testGroup "Hessenberg" + [ testProperty "A = QHQᵀ reconstruction" prop_hessenbergReconstruction + , testProperty "H is upper Hessenberg" prop_hessenbergForm + ] + , testGroup "Symmetric" + [ testProperty "A = QΛQᵀ reconstruction" prop_symmetricEigenReconstruction + , testProperty "Q orthogonal" prop_symmetricQOrthogonal + , testCase "eigenvalues of diagonal" test_symmetricDiagonal + ] + , testGroup "Jacobi" + [ testCase "jacobi eigenvalues of known matrix" test_jacobiKnown + ] + , testGroup "SVD" + [ testProperty "A ≈ UΣVᵀ reconstruction" prop_svdReconstruction + , testCase "singular values of diagonal" test_svdDiagonal + , testCase "svdGKP reconstruction 10x10" test_svdGKReconstruction + ] + , testGroup "Standard test matrices" + [ testCase "Wilkinson eigenvalues" test_wilkinsonEigen + , testCase "Hilbert eigenvalues positive" test_hilbertEigen + , testCase "Frank eigenvalues positive" test_frankEigen + , testProperty "clustered eigenvalues" prop_clusteredEigen + ] + , testGroup "Eigen residuals" + [ testProperty "eigenpair scaled residuals 3x3" prop_eigenResiduals + ] + , testGroup "SVD residuals" + [ testProperty "SVD scaled residual 3x3" prop_svdScaledResidual + , testProperty "SVD orthogonality U and V 3x3" prop_svdOrthogonality + , testCase "SVD diagonal 5x5 sorted" test_svdDiagonalLarger + ] + , testGroup "Eigenvalue ordering" + [ testProperty "symmetric eigenvalues sorted 4x4" prop_symmetricEigenvaluesSorted + ] + , testGroup "D&C eigensolver" + [ testCase "D&C eigenvalues of diagonal 10x10" test_dcEigenDiagonal + , testCase "D&C reconstruction 50x50" test_dcEigenReconstruction50 + , testCase "D&C orthogonality 50x50" test_dcEigenOrthogonal50 + , testCase "D&C matches QR at 30x30" test_dcMatchesQR + , testCase "D&C orthogonality 30x30" test_dcOrtho30 + , testCase "D&C orthogonality 52x52" test_dcOrtho52 + , testCase "D&C orthogonality 60x60" test_dcOrtho60 + , testCase "D&C orthogonality 80x80" test_dcOrtho80 + , testCase "D&C orthogonality 90x90" test_dcOrtho90 + , testCase "D&C orthogonality 95x95" test_dcOrtho95 + , testCase "D&C ortho diagonal 100x100" test_dcOrthoDiag100 + , testCase "D&C ortho alt-matrix 100x100" test_dcOrthoAlt100 + , testCase "D&C reconstruction 100x100" test_dcEigenReconstruction100 + , testCase "D&C reconstruction 128x128" test_dcEigenReconstruction128 + ] + , testGroup "Panel tridiag (n >= 256)" + [ testCase "tridiag match 128x128" test_panelTridiagReconstruction128 + , testCase "eigenreconstruction 200x200" test_panelTridiagReconstruction200 + , testCase "orthogonality 200x200" test_panelTridiagOrthogonal200 + , testCase "eigenreconstruction 300x300" test_panelTridiagReconstruction300 + ] + ] + +-- Power method + +test_powerDiagonal :: Assertion +test_powerDiagonal = do + -- A = diag(3, 2, 1) → dominant eigenvalue = 3 + let a = makeMatrix @3 @3 @M.P $ \i j -> + if i == j then case i of { 0 -> 3; 1 -> 2; _ -> 1 } else 0 :: Double + q0 = makeVector @3 @M.P $ \_ -> 1 / sqrt 3 :: Double + (lam, _) = powerMethod a q0 100 1e-10 + assertBool "eigenvalue ~ 3" $ abs (lam - 3) < 0.01 + +-- Hessenberg + +prop_hessenbergReconstruction :: Property +prop_hessenbergReconstruction = forAll (genMatrix @4 @4) $ \a -> + let (q, h) = hessenberg a + qt = transpose q + qhqt = matMul q (matMul h qt) + in matApproxEq @4 @4 a qhqt + +prop_hessenbergForm :: Property +prop_hessenbergForm = forAll (genMatrix @4 @4) $ \a -> + let (_, h) = hessenberg a + in all (\(i, j) -> abs (h ! (i, j)) < 1e-8) + [(i, j) | i <- [0..3], j <- [0..3], i > j + 1] + +-- Symmetric eigenvalue + +prop_symmetricEigenReconstruction :: Property +prop_symmetricEigenReconstruction = forAll (genSPDMatrix @3) $ \a -> + let (eigvals, q) = symmetricEigen a 500 1e-12 + qt = transpose q + -- Reconstruct: A ≈ Q diag(λ) Qᵀ + lambda = makeMatrix @3 @3 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul lambda qt) + -- Use relaxed tolerance for iterative eigenvalue decomposition + in all (\(i, j) -> abs (a ! (i,j) - qlqt ! (i,j)) < 1e-4 * (1 + abs (a ! (i,j)))) + [(i, j) | i <- [0..2], j <- [0..2]] + +prop_symmetricQOrthogonal :: Property +prop_symmetricQOrthogonal = forAll (genSPDMatrix @3) $ \a -> + let (_, q) = symmetricEigen a 500 1e-12 + qt = transpose q + qtq = matMul qt q + eye = identityMatrix @3 @M.P :: Matrix 3 3 M.P Double + in matApproxEq @3 @3 qtq eye + +test_symmetricDiagonal :: Assertion +test_symmetricDiagonal = do + -- Eigenvalues of diag(5, 3, 1) should be {1, 3, 5} + let a = makeMatrix @3 @3 @M.P $ \i j -> + if i == j then case i of { 0 -> 5; 1 -> 3; _ -> 1 } else 0 :: Double + (eigvals, _) = symmetricEigen a 100 1e-12 + evs = sort [eigvals !. 0, eigvals !. 1, eigvals !. 2] + assertBool "eigenvalue 1" $ abs (evs !! 0 - 1) < 0.01 + assertBool "eigenvalue 3" $ abs (evs !! 1 - 3) < 0.01 + assertBool "eigenvalue 5" $ abs (evs !! 2 - 5) < 0.01 + +-- Jacobi + +test_jacobiKnown :: Assertion +test_jacobiKnown = do + let a = makeMatrix @3 @3 @M.P $ \i j -> case (i,j) of + (0,0) -> 4; (0,1) -> 1; (0,2) -> 0 + (1,0) -> 1; (1,1) -> 3; (1,2) -> 1 + (2,0) -> 0; (2,1) -> 1; (2,2) -> 2 + _ -> 0 :: Double + (eigvals, q) = jacobiEigen a 100 1e-12 + qt = transpose q + lambda = makeMatrix @3 @3 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul lambda qt) + assertBool "Jacobi reconstruction" $ matApproxEq @3 @3 a qlqt + +-- SVD + +prop_svdReconstruction :: Property +prop_svdReconstruction = forAll (genMatrix @3 @3) $ \a -> + let (u, sigma, v) = svd a + vt = transpose v + -- Reconstruct: U * diag(σ) * Vᵀ + sigMat = makeMatrix @3 @3 @M.P $ \i j -> + if i == j then sigma !. i else 0 + usv = matMul u (matMul sigMat vt) + in matApproxEq @3 @3 a usv + +test_svdDiagonal :: Assertion +test_svdDiagonal = do + -- SVD of diag(5, 3, 1) should give singular values {5, 3, 1} + let a = makeMatrix @3 @3 @M.P $ \i j -> + if i == j then case i of { 0 -> 5; 1 -> 3; _ -> 1 } else 0 :: Double + (_, sigma, _) = svd a + svs = sort [sigma !. 0, sigma !. 1, sigma !. 2] + assertBool "sv 1" $ abs (svs !! 0 - 1) < 0.1 + assertBool "sv 3" $ abs (svs !! 1 - 3) < 0.1 + assertBool "sv 5" $ abs (svs !! 2 - 5) < 0.1 + +test_svdGKReconstruction :: Assertion +test_svdGKReconstruction = do + -- Test 1: diagonal matrix (trivial bidiag, no QR needed) + let diag5 = makeMatrix @5 @5 @M.P $ \i j -> + if i == j then fromIntegral (5 - i) else 0 :: Double + (_, sigDiag, _) = svdGKP diag5 + diagSorted = sort [sigDiag !. i | i <- [0..4]] + diagExpected = [1,2,3,4,5] :: [Double] + diagErr = maximum $ zipWith (\a_ b_ -> abs (a_ - b_)) diagSorted diagExpected + assertBool ("svdGKP diagonal sigma " ++ show diagSorted ++ " err=" ++ show diagErr) $ diagErr < 0.1 + -- Test 2: already-bidiagonal matrix (tests QR iteration in isolation) + -- B = [[3,1,0],[0,2,1],[0,0,1]] — bidiag with d=[3,2,1], e=[1,1] + let bidiag3 = makeMatrix @3 @3 @M.P $ \i j -> case (i,j) of + (0,0) -> 3; (0,1) -> 1; (1,1) -> 2; (1,2) -> 1; (2,2) -> 1 + _ -> 0 :: Double + (_, sigBidiag, _) = svdGKP bidiag3 + (_, sigBidiagRef, _) = svdP bidiag3 + bidiagSorted = sort [sigBidiag !. i | i <- [0..2]] + bidiagRefSorted = sort [sigBidiagRef !. i | i <- [0..2]] + bidiagDiff = maximum $ zipWith (\a_ b_ -> abs (a_ - b_)) bidiagSorted bidiagRefSorted + assertBool ("svdGKP bidiag diff " ++ show bidiagDiff + ++ "\n gk=" ++ show bidiagSorted + ++ "\n ref=" ++ show bidiagRefSorted) $ bidiagDiff < 1e-6 + -- Test 3: general matrix + let a = makeMatrix @5 @5 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in 1.0 / d + if i == j then 5 else 0 + (_, sigmaGK, _) = svdGKP a + (_, sigmaRef, _) = svdP a + gkSorted = sort [sigmaGK !. i | i <- [0..4]] + refSorted = sort [sigmaRef !. i | i <- [0..4]] + svDiff = maximum $ zipWith (\a_ b_ -> abs (a_ - b_)) gkSorted refSorted + assertBool ("svdGKP sigma diff " ++ show svDiff ++ "\n gk=" ++ show gkSorted + ++ "\n ref=" ++ show refSorted) $ svDiff < 1e-6 + -- Test 4: singular values match for 10×10 + let a10 = makeMatrix @10 @10 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in 1.0 / d + if i == j then 10 else 0 + (u10, sig10, v10) = svdGKP a10 + (_, sigRef10, _) = svdP a10 + gk10Sorted = sort [sig10 !. i | i <- [0..9]] + ref10Sorted = sort [sigRef10 !. i | i <- [0..9]] + svDiff10 = maximum $ zipWith (\a_ b_ -> abs (a_ - b_)) gk10Sorted ref10Sorted + assertBool ("svdGKP 10x10 sigma diff " ++ show svDiff10 + ++ "\n gk=" ++ show gk10Sorted + ++ "\n ref=" ++ show ref10Sorted) $ svDiff10 < 1e-6 + -- Test 5: reconstruction A ≈ U Σ V^T for 10×10 + let sigMat10 = makeMatrix @10 @10 @M.P $ \i j -> + if i == j then sig10 !. i else 0 + usv10 = matMulP u10 (matMulP sigMat10 (transpose v10)) + reconErr10 = maximum [abs (a10 ! (i,j) - usv10 ! (i,j)) + | i <- [0..9], j <- [0..9]] + assertBool ("svdGKP 10x10 reconstruction err " ++ show reconErr10) $ reconErr10 < 1e-10 + -- Test 6: orthogonality of U and V + let utu = matMulP (transpose u10) u10 + vtv = matMulP (transpose v10) v10 + eye10 = identityMatrix @10 @M.P :: Matrix 10 10 M.P Double + uErr = maximum [abs (utu ! (i,j) - eye10 ! (i,j)) | i <- [0..9], j <- [0..9]] + vErr = maximum [abs (vtv ! (i,j) - eye10 ! (i,j)) | i <- [0..9], j <- [0..9]] + assertBool ("svdGKP 10x10 U ortho err " ++ show uErr) $ uErr < 1e-10 + assertBool ("svdGKP 10x10 V ortho err " ++ show vErr) $ vErr < 1e-10 + +-- Standard test matrices + +test_wilkinsonEigen :: Assertion +test_wilkinsonEigen = do + let a = wilkinsonMatrix @7 :: Matrix 7 7 M.P Double + (eigvals, q) = symmetricEigen a 500 1e-12 + nn = 7 + -- Verify we get 7 eigenvalues + evList = map (\i -> eigvals !. i) [0..nn-1] + assertBool "got 7 eigenvalues" $ length evList == 7 + -- Verify reconstruction: A ≈ Q diag(λ) Qᵀ + let diag_lambda = makeMatrix @7 @7 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qt = transpose q + qlqt = matMul q (matMul diag_lambda qt) + residual = normFrob (mSub a qlqt) / (normFrob a + 1e-15) + assertBool "Wilkinson reconstruction" $ residual < 1e-4 + +test_hilbertEigen :: Assertion +test_hilbertEigen = do + let a = hilbertMatrix @5 :: Matrix 5 5 M.P Double + (eigvals, _) = symmetricEigen a 500 1e-12 + evList = map (\i -> eigvals !. i) [0..4] + -- Hilbert matrix is SPD, so all eigenvalues must be positive + assertBool "all eigenvalues positive" $ all (> 0) evList + +test_frankEigen :: Assertion +test_frankEigen = do + let a = frankMatrix @5 :: Matrix 5 5 M.P Double + (_, t) = schur a 200 1e-10 + evs = eigenvalues @5 t + -- Frank matrix has all positive real eigenvalues + assertBool "all eigenvalues positive" $ all (> 0) evs + +prop_clusteredEigen :: Property +prop_clusteredEigen = withMaxSuccess 10 $ forAll (genClusteredEigenMatrix @4 5.0) $ \a -> + let (eigvals, q) = symmetricEigen a 500 1e-12 + qt = transpose q + diag_lambda = makeMatrix @4 @4 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul diag_lambda qt) + residual = normFrob (mSub a qlqt) / (normFrob a + 1e-15) + -- Relaxed tolerance since clustered eigenvalues are harder + in residual < 1e-3 + +-- Eigen residuals + +prop_eigenResiduals :: Property +prop_eigenResiduals = withMaxSuccess 20 $ forAll (genSPDMatrix @3) $ \a -> + let (eigvals, q) = symmetricEigen a 500 1e-12 + nn = 3 + -- Check each eigenpair + checks = map (\i -> + let lambda_i = eigvals !. i + v_i = makeVector @3 @M.P $ \k -> q ! (k, i) + in scaledResidualEigen a lambda_i v_i < 1000 + ) [0..nn-1] + in all id checks + +-- SVD residuals + +prop_svdScaledResidual :: Property +prop_svdScaledResidual = forAll (genMatrix @3 @3) $ \a -> + let (u, sigma, v) = svd a + in scaledResidualSVD a u sigma v < 1000 + +prop_svdOrthogonality :: Property +prop_svdOrthogonality = forAll (genMatrix @3 @3) $ \a -> + let (u, _, v) = svd a + -- U orthogonality can be looser because the SVD implementation + -- constructs U columns as Av/sigma, which may accumulate error. + -- V comes from eigendecomposition of A^T A so is typically tighter. + in orthogonalityResidual @3 u < 500000 && orthogonalityResidual @3 v < 5000 + +test_svdDiagonalLarger :: Assertion +test_svdDiagonalLarger = do + let a = makeMatrix @5 @5 @M.P $ \i j -> + if i == j then case i of + 0 -> 7; 1 -> 5; 2 -> 3; 3 -> 2; _ -> 1 + else 0 :: Double + (_, sigma, _) = svd a + svs = sort [sigma !. 0, sigma !. 1, sigma !. 2, sigma !. 3, sigma !. 4] + expected = [1, 2, 3, 5, 7] :: [Double] + assertBool "sorted singular values match" $ + all (\(s, e) -> abs (s - e) < 0.1) (zip svs expected) + +-- Eigenvalue ordering + +prop_symmetricEigenvaluesSorted :: Property +prop_symmetricEigenvaluesSorted = forAll (genSPDMatrix @4) $ \a -> + let (eigvals, _) = symmetricEigen a 500 1e-12 + evList = sort $ map (\i -> eigvals !. i) [0..3] + -- Verify non-decreasing order after sorting + in and $ zipWith (<=) evList (tail evList) + +-- D&C eigensolver tests + +test_dcEigenDiagonal :: Assertion +test_dcEigenDiagonal = do + -- Eigenvalues of diag(10, 9, 8, ..., 1) should be {1..10} + let a = makeMatrix @10 @10 @M.P $ \i j -> + if i == j then fromIntegral (10 - i) else 0 :: Double + (eigvals, _) = symmetricEigenPDC a 1e-12 + evs = sort [eigvals !. i | i <- [0..9]] + mapM_ (\(i, expected) -> + assertBool ("eigenvalue " ++ show expected) $ + abs (evs !! i - expected) < 0.01) + (zip [0..] [1..10 :: Double]) + +test_dcEigenReconstruction50 :: Assertion +test_dcEigenReconstruction50 = do + -- A = QΛQ^T reconstruction for a 50x50 SPD matrix + let a = mkSPD50 + (eigvals, q) = symmetricEigenPDC a 1e-12 + qt = transpose q + lambda = makeMatrix @50 @50 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul lambda qt) + maxErr = maximum [abs (a ! (i,j) - qlqt ! (i,j)) | i <- [0..49], j <- [0..49]] + assertBool ("reconstruction error " ++ show maxErr ++ " < 1e-8") $ maxErr < 1e-8 + +test_dcEigenOrthogonal50 :: Assertion +test_dcEigenOrthogonal50 = do + let a = mkSPD50 + (_, q) = symmetricEigenPDC a 1e-12 + qt = transpose q + qtq = matMul qt q + eye = identityMatrix @50 @M.P :: Matrix 50 50 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..49], j <- [0..49]] + assertBool ("orthogonality error " ++ show maxErr ++ " < 1e-8") $ maxErr < 1e-8 + +test_dcMatchesQR :: Assertion +test_dcMatchesQR = do + -- D&C and QR should produce same eigenvalues for a 30x30 SPD matrix + let a = makeMatrix @30 @30 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 30 + fromIntegral i else 1.0 / d + (eigsDC, _) = symmetricEigenPDC a 1e-12 + (eigsQR, _) = symmetricEigenP a 3000 1e-12 + dcSorted = sort [eigsDC !. i | i <- [0..29]] + qrSorted = sort [eigsQR !. i | i <- [0..29]] + maxDiff = maximum $ zipWith (\a' b' -> abs (a' - b')) dcSorted qrSorted + assertBool ("D&C vs QR diff " ++ show maxDiff ++ " < 1e-8") $ maxDiff < 1e-8 + +test_dcEigenReconstruction100 :: Assertion +test_dcEigenReconstruction100 = do + let a = makeMatrix @100 @100 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 100 + fromIntegral i else 1.0 / d + (eigvalsDC, qDC) = symmetricEigenPDC a 1e-12 + (eigvalsQR, _) = symmetricEigenP a 10000 1e-12 + -- Compare eigenvalues + dcSorted = sort [eigvalsDC !. i | i <- [0..99]] + qrSorted = sort [eigvalsQR !. i | i <- [0..99]] + evDiff = maximum $ zipWith (\a_ b_ -> abs (a_ - b_)) dcSorted qrSorted + assertBool ("D&C 100 eigenvalue diff " ++ show evDiff) $ evDiff < 1e-6 + -- Check orthogonality of Q + let qtq = matMulP (transpose qDC) qDC + eye100 = identityMatrix @100 @M.P :: Matrix 100 100 M.P Double + orthoErr = maximum [abs (qtq ! (i,j) - eye100 ! (i,j)) | i <- [0..99], j <- [0..99]] + assertBool ("D&C 100 orthogonality error " ++ show orthoErr) $ orthoErr < 1e-6 + -- Full reconstruction + let qt = transpose qDC + lambda = makeMatrix @100 @100 @M.P $ \i j -> + if i == j then eigvalsDC !. i else 0 + qlqt = matMul qDC (matMul lambda qt) + maxErr = maximum [abs (a ! (i,j) - qlqt ! (i,j)) | i <- [0..99], j <- [0..99]] + assertBool ("D&C 100 reconstruction error " ++ show maxErr ++ " < 1e-7") $ maxErr < 1e-7 + +test_dcEigenReconstruction128 :: Assertion +test_dcEigenReconstruction128 = do + let a = makeMatrix @128 @128 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 128 + fromIntegral i else 1.0 / d + (eigvals, q) = symmetricEigenPDC a 1e-12 + qt = transpose q + lambda = makeMatrix @128 @128 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul lambda qt) + maxErr = maximum [abs (a ! (i,j) - qlqt ! (i,j)) | i <- [0..127], j <- [0..127]] + assertBool ("D&C 128 reconstruction error " ++ show maxErr ++ " < 5e-7") $ maxErr < 5e-7 + +-- Panel tridiag tests (n >= 128 crossover) + +test_panelTridiagReconstruction128 :: Assertion +test_panelTridiagReconstruction128 = do + let nn = 128 + a = makeMatrix @128 @128 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 128 + fromIntegral i else 1.0 / d + (eigvals, q) = symmetricEigenP a 10000 1e-12 + qt = transpose q + lambda = makeMatrix @128 @128 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul lambda qt) + maxErr = maximum [abs (a ! (i,j) - qlqt ! (i,j)) | i <- [0..nn-1], j <- [0..nn-1]] + assertBool ("reconstruction error " ++ show maxErr ++ " < 1e-6") $ maxErr < 1e-6 + +test_panelTridiagReconstruction200 :: Assertion +test_panelTridiagReconstruction200 = do + let nn = 200 + a = mkSPD200 + (eigvals, q) = symmetricEigenP a 10000 1e-12 + qt = transpose q + lambda = makeMatrix @200 @200 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul lambda qt) + maxErr = maximum [abs (a ! (i,j) - qlqt ! (i,j)) | i <- [0..nn-1], j <- [0..nn-1]] + assertBool ("reconstruction error " ++ show maxErr ++ " < 1e-6") $ maxErr < 1e-6 + +test_panelTridiagOrthogonal200 :: Assertion +test_panelTridiagOrthogonal200 = do + let nn = 200 + a = mkSPD200 + (_, q) = symmetricEigenP a 10000 1e-12 + qt = transpose q + qtq = matMul qt q + eye = identityMatrix @200 @M.P :: Matrix 200 200 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..nn-1], j <- [0..nn-1]] + assertBool ("orthogonality error " ++ show maxErr ++ " < 1e-8") $ maxErr < 1e-8 + +test_panelTridiagReconstruction300 :: Assertion +test_panelTridiagReconstruction300 = do + let nn = 300 + a = makeMatrix @300 @300 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 300 + fromIntegral i else 1.0 / d + (eigvals, q) = symmetricEigenP a 15000 1e-12 + qt = transpose q + lambda = makeMatrix @300 @300 @M.P $ \i j -> + if i == j then eigvals !. i else 0 + qlqt = matMul q (matMul lambda qt) + maxErr = maximum [abs (a ! (i,j) - qlqt ! (i,j)) | i <- [0..nn-1], j <- [0..nn-1]] + assertBool ("reconstruction error " ++ show maxErr ++ " < 1e-6") $ maxErr < 1e-6 + +mkSPD200 :: Matrix 200 200 M.P Double +mkSPD200 = makeMatrix @200 @200 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 200 + fromIntegral i else 1.0 / d + +-- Helper: 50x50 SPD matrix for D&C tests +mkSPD50 :: Matrix 50 50 M.P Double +mkSPD50 = makeMatrix @50 @50 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 50 + fromIntegral i else 1.0 / d + +-- Granular D&C orthogonality tests at various sizes +test_dcOrtho30 :: Assertion +test_dcOrtho30 = do + let a = makeMatrix @30 @30 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 30 + fromIntegral i else 1.0 / d + (_, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @30 @M.P :: Matrix 30 30 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..29], j <- [0..29]] + assertBool ("D&C 30 ortho err " ++ show maxErr) $ maxErr < 1e-8 + +test_dcOrtho52 :: Assertion +test_dcOrtho52 = do + let a = makeMatrix @52 @52 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 52 + fromIntegral i else 1.0 / d + (_, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @52 @M.P :: Matrix 52 52 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..51], j <- [0..51]] + assertBool ("D&C 52 ortho err " ++ show maxErr) $ maxErr < 1e-8 + +test_dcOrtho60 :: Assertion +test_dcOrtho60 = do + let a = makeMatrix @60 @60 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 60 + fromIntegral i else 1.0 / d + (_, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @60 @M.P :: Matrix 60 60 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..59], j <- [0..59]] + assertBool ("D&C 60 ortho err " ++ show maxErr) $ maxErr < 1e-8 + +test_dcOrtho80 :: Assertion +test_dcOrtho80 = do + let a = makeMatrix @80 @80 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 80 + fromIntegral i else 1.0 / d + (_, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @80 @M.P :: Matrix 80 80 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..79], j <- [0..79]] + assertBool ("D&C 80 ortho err " ++ show maxErr) $ maxErr < 1e-8 + +test_dcOrtho90 :: Assertion +test_dcOrtho90 = do + let a = makeMatrix @90 @90 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 90 + fromIntegral i else 1.0 / d + (_, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @90 @M.P :: Matrix 90 90 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..89], j <- [0..89]] + assertBool ("D&C 90 ortho err " ++ show maxErr) $ maxErr < 1e-8 + +test_dcOrtho95 :: Assertion +test_dcOrtho95 = do + let a = makeMatrix @95 @95 @M.P $ \i j -> + let d = fromIntegral (abs (i - j) + 1) :: Double + in if i == j then 95 + fromIntegral i else 1.0 / d + (_, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @95 @M.P :: Matrix 95 95 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..94], j <- [0..94]] + assertBool ("D&C 95 ortho err " ++ show maxErr) $ maxErr < 1e-8 + +-- Test D&C with a purely diagonal 100×100 matrix +test_dcOrthoDiag100 :: Assertion +test_dcOrthoDiag100 = do + let a = makeMatrix @100 @100 @M.P $ \i j -> + if i == j then fromIntegral (i + 1) else 0 :: Double + (eigvals, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @100 @M.P :: Matrix 100 100 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..99], j <- [0..99]] + evSorted = sort [eigvals !. i | i <- [0..99]] + evDiff = maximum $ zipWith (\a_ b_ -> abs (a_ - b_)) evSorted [1..100] + assertBool ("D&C diag100 ortho err " ++ show maxErr) $ maxErr < 1e-8 + assertBool ("D&C diag100 eigenvalue diff " ++ show evDiff) $ evDiff < 1e-8 + +-- Test D&C with a different matrix at 100×100 (sparser off-diagonal) +test_dcOrthoAlt100 :: Assertion +test_dcOrthoAlt100 = do + let a = makeMatrix @100 @100 @M.P $ \i j -> + if i == j then 500 + fromIntegral i + else if abs (i - j) == 1 then 0.1 + else 0 :: Double + (_, q) = symmetricEigenPDC a 1e-12 + qtq = matMulP (transpose q) q + eye = identityMatrix @100 @M.P :: Matrix 100 100 M.P Double + maxErr = maximum [abs (qtq ! (i,j) - eye ! (i,j)) | i <- [0..99], j <- [0..99]] + assertBool ("D&C alt100 ortho err " ++ show maxErr) $ maxErr < 1e-8 diff --git a/test/Test/Norms.hs b/test/Test/Norms.hs new file mode 100644 --- /dev/null +++ b/test/Test/Norms.hs @@ -0,0 +1,112 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Test.Norms (normTests) where + +import Test.Tasty +import Test.Tasty.QuickCheck +import Test.Tasty.HUnit +import qualified Data.Massiv.Array as M + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.Norms +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (transpose, matMul) +import Numeric.LinearAlgebra.Massiv.Eigen.SVD (singularValues) +import Test.Types + +normTests :: TestTree +normTests = testGroup "Norms" + [ testGroup "Vector norms" + [ testProperty "vnorm2 non-negative" prop_vnorm2NonNeg + , testProperty "vnorm1 non-negative" prop_vnorm1NonNeg + , testProperty "vnormInf <= vnorm1" prop_vnormInfLeqVnorm1 + , testCase "known vector norm" test_knownVnorm + ] + , testGroup "Matrix norms" + [ testProperty "normFrob non-negative" prop_normFrobNonNeg + , testProperty "norm1 = normInf of transpose" prop_norm1TransposeNormInf + , testProperty "triangle inequality (Frobenius)" prop_triangleInequality + , testProperty "norm consistency with sigma_max" prop_normConsistency + , testProperty "submultiplicativity" prop_submultiplicativity + , testCase "known Frobenius norm" test_knownFrobNorm + , testCase "Hilbert matrix norm relationships" test_hilbertNorms + ] + ] + +-- Vector norms + +prop_vnorm2NonNeg :: Property +prop_vnorm2NonNeg = forAll (genVector @4) $ \x -> + vnorm2 x >= (0 :: Double) + +prop_vnorm1NonNeg :: Property +prop_vnorm1NonNeg = forAll (genVector @4) $ \x -> + vnorm1 x >= (0 :: Double) + +prop_vnormInfLeqVnorm1 :: Property +prop_vnormInfLeqVnorm1 = forAll (genVector @4) $ \x -> + vnormInf x <= vnorm1 x + 1e-12 + +test_knownVnorm :: Assertion +test_knownVnorm = do + let v = makeVector @3 @M.P $ \i -> case i of { 0 -> 3; 1 -> 4; _ -> 0 } :: Double + assertBool "vnorm2 of [3,4,0] = 5" $ abs (vnorm2 v - 5) < 1e-12 + assertBool "vnorm1 of [3,4,0] = 7" $ abs (vnorm1 v - 7) < 1e-12 + assertBool "vnormInf of [3,4,0] = 4" $ abs (vnormInf v - 4) < 1e-12 + +-- Matrix norms + +prop_normFrobNonNeg :: Property +prop_normFrobNonNeg = forAll (genMatrix @3 @3) $ \a -> + normFrob a >= (0 :: Double) + +prop_norm1TransposeNormInf :: Property +prop_norm1TransposeNormInf = forAll (genMatrix @3 @4) $ \a -> + let at = transpose a + in abs (norm1 a - normInf at) < 1e-8 + +prop_triangleInequality :: Property +prop_triangleInequality = forAll ((,) <$> genMatrix @3 @3 <*> genMatrix @3 @3) $ \(a, b) -> + let ab = makeMatrix @3 @3 @M.P $ \i j -> (a ! (i,j)) + (b ! (i,j)) + in normFrob ab <= normFrob a + normFrob b + 1e-10 + +-- | For 5×5 matrices: sigma_max ≤ normFrob A ≤ sqrt(5) * sigma_max. +-- +-- The Frobenius norm satisfies ‖A‖_F = sqrt(sum sigma_i^2), so +-- sigma_max ≤ ‖A‖_F ≤ sqrt(n) * sigma_max. +prop_normConsistency :: Property +prop_normConsistency = forAll (genMatrix @5 @5) $ \a -> + let sv = singularValues a + sigmaMax = sv !. 0 + frobA = normFrob a + in sigmaMax <= frobA + 1e-10 + && frobA <= sqrt 5 * sigmaMax + 1e-10 + +-- | Submultiplicativity: ‖A·B‖_F ≤ ‖A‖_F · ‖B‖_F for 3×3 matrices. +prop_submultiplicativity :: Property +prop_submultiplicativity = + forAll ((,) <$> genMatrix @3 @3 <*> genMatrix @3 @3) $ \(a, b) -> + normFrob (matMul a b) <= normFrob a * normFrob b + 1e-10 + +-- | Hilbert matrix norm relationships for hilbertMatrix @4: +-- norm1 ≤ normFrob * sqrt(n) and normFrob ≤ sqrt(n) * normInf. +test_hilbertNorms :: Assertion +test_hilbertNorms = do + let h = hilbertMatrix @4 :: Matrix 4 4 M.P Double + frobH = normFrob h + n1H = norm1 h + niH = normInf h + sqrtN = sqrt 4 :: Double + assertBool ("norm1 <= normFrob * sqrt(4): norm1=" ++ show n1H + ++ " normFrob*2=" ++ show (frobH * sqrtN)) + $ n1H <= frobH * sqrtN + 1e-12 + assertBool ("normFrob <= sqrt(4) * normInf: normFrob=" ++ show frobH + ++ " 2*normInf=" ++ show (sqrtN * niH)) + $ frobH <= sqrtN * niH + 1e-12 + +test_knownFrobNorm :: Assertion +test_knownFrobNorm = do + -- ‖[[1,2],[3,4]]‖_F = sqrt(1+4+9+16) = sqrt(30) + let a = makeMatrix @2 @2 @M.P $ \i j -> case (i,j) of + (0,0) -> 1; (0,1) -> 2; (1,0) -> 3; (1,1) -> 4; _ -> 0 :: Double + assertBool "Frobenius norm" $ abs (normFrob a - sqrt 30) < 1e-12 diff --git a/test/Test/Orthogonal.hs b/test/Test/Orthogonal.hs new file mode 100644 --- /dev/null +++ b/test/Test/Orthogonal.hs @@ -0,0 +1,198 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Test.Orthogonal (orthogonalTests) where + +import Test.Tasty +import Test.Tasty.QuickCheck +import Test.Tasty.HUnit +import qualified Data.Massiv.Array as M + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, transpose, mSub) +import Numeric.LinearAlgebra.Massiv.Orthogonal.Householder +import Numeric.LinearAlgebra.Massiv.Orthogonal.Givens +import Numeric.LinearAlgebra.Massiv.Orthogonal.QR (qr, qrGivens) +import Numeric.LinearAlgebra.Massiv.Orthogonal.LeastSquares +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvec) +import Numeric.LinearAlgebra.Massiv.Norms (normFrob, vnorm2) +import Test.Types +import Test.Residuals + +orthogonalTests :: TestTree +orthogonalTests = testGroup "Orthogonal" + [ testGroup "Householder" + [ testProperty "Householder zeros subdiagonal" prop_householderZeros + , testProperty "Householder matrix orthogonal" prop_householderOrthogonal + ] + , testGroup "Givens" + [ testProperty "Givens zeros target" prop_givensZeros + , testProperty "Givens preserves norm" prop_givensNorm + ] + , testGroup "QR" + [ testProperty "A = QR reconstruction" prop_qrReconstruction + , testProperty "Q orthogonal" prop_qOrthogonal + , testProperty "R upper triangular" prop_rUpperTriangular + , testCase "3x3 QR known" test_qrKnown + ] + , testGroup "Least Squares" + [ testCase "overdetermined system" test_leastSquares + ] + , testGroup "QR scaled residuals" + [ testProperty "QR scaled residual 5x5" prop_qrScaledResidual5 + , testProperty "Q orthogonality residual 5x5" prop_qOrthogonalityResidual5 + , testProperty "Q orthogonality residual 10x10" prop_qOrthogonalityResidual10 + ] + , testGroup "QR reconstruction (larger)" + [ testProperty "QR reconstruction 10x10" prop_qrReconstruction10x10 + , testProperty "QR reconstruction 10x5 rectangular" prop_qrReconstruction10x5 + ] + , testGroup "Givens vs Householder" + [ testProperty "Givens matches Householder 5x5" prop_givensMatchesHouseholder + ] + , testGroup "Least squares (property)" + [ testProperty "normal equations 10x5" prop_leastSquares10x5 + ] + ] + +-- Householder tests + +prop_householderZeros :: Property +prop_householderZeros = forAll (genVector @4) $ \x -> + let (v, beta) = householderVector x + h = householderMatrix @4 @M.P v beta + hx = matvec h x + -- All entries below first should be ~0 + in all (\i -> abs (hx !. i) < 1e-6) [1..3] + +prop_householderOrthogonal :: Property +prop_householderOrthogonal = forAll (genVector @3) $ \x -> + let (v, beta) = householderVector x + h = householderMatrix @3 @M.P v beta + ht = transpose h + hht = matMul h ht + eye = identityMatrix @3 @M.P :: Matrix 3 3 M.P Double + in matApproxEq @3 @3 hht eye + +-- Givens tests + +prop_givensZeros :: Property +prop_givensZeros = forAll ((,) <$> choose (-10, 10) <*> choose (-10, 10)) $ \(a, b) -> + let (c, s) = givensRotation a b + -- Convention: [c, -s; s, c] * [a; b] = [r; 0] + zero = s * a + c * b -- should be ~0 + in abs zero < (1e-10 :: Double) + +prop_givensNorm :: Property +prop_givensNorm = forAll ((,) <$> choose (-10, 10) <*> choose (-10, 10)) $ \(a, b) -> + let (c, s) = givensRotation a (b :: Double) + r = c * a - s * b -- r = sqrt(a² + b²) + in abs (r * r - (a * a + b * b)) < 1e-8 + +-- QR tests + +prop_qrReconstruction :: Property +prop_qrReconstruction = forAll (genMatrix @4 @3) $ \a -> + let (q, r) = qr a + qr_ = matMul q r + in matApproxEq @4 @3 a qr_ + +prop_qOrthogonal :: Property +prop_qOrthogonal = forAll (genMatrix @3 @3) $ \a -> + let (q, _) = qr a + qt = transpose q + qtq = matMul qt q + eye = identityMatrix @3 @M.P :: Matrix 3 3 M.P Double + in matApproxEq @3 @3 qtq eye + +prop_rUpperTriangular :: Property +prop_rUpperTriangular = forAll (genMatrix @3 @3) $ \a -> + let (_, r) = qr a + in all (\(i, j) -> abs (r ! (i, j)) < 1e-8) + [(i, j) | i <- [0..2], j <- [0..2], i > j] + +test_qrKnown :: Assertion +test_qrKnown = do + let a = makeMatrix @3 @3 @M.P $ \i j -> case (i,j) of + (0,0) -> 12; (0,1) -> -51; (0,2) -> 4 + (1,0) -> 6; (1,1) -> 167; (1,2) -> -68 + (2,0) -> -4; (2,1) -> 24; (2,2) -> -41 + _ -> 0 :: Double + (q, r) = qr a + qr_ = matMul q r + assertBool "QR reconstruction" $ matApproxEq @3 @3 a qr_ + +-- Least squares test + +test_leastSquares :: Assertion +test_leastSquares = do + -- Overdetermined system: 4 equations, 2 unknowns + -- A = [[1,1],[1,2],[1,3],[1,4]], b = [6,5,7,10] + -- Least squares fit: y = a + bx + let a = makeMatrix @4 @2 @M.P $ \i j -> case (i,j) of + (0,0) -> 1; (0,1) -> 1 + (1,0) -> 1; (1,1) -> 2 + (2,0) -> 1; (2,1) -> 3 + (3,0) -> 1; (3,1) -> 4 + _ -> 0 :: Double + b = makeVector @4 @M.P $ \i -> case i of + 0 -> 6; 1 -> 5; 2 -> 7; _ -> 10 :: Double + x = leastSquaresQR a b + -- The solution should minimize ‖Ax - b‖₂ + -- With these values: x ≈ [3.5, 1.4] + assertBool "intercept reasonable" $ abs (x !. 0 - 3.5) < 0.5 + assertBool "slope reasonable" $ abs (x !. 1 - 1.4) < 0.5 + +-- QR scaled residual tests + +prop_qrScaledResidual5 :: Property +prop_qrScaledResidual5 = forAll (genMatrix @5 @5) $ \a -> + let (q, r) = qr a + in scaledResidualQR a q r < 100 + +prop_qOrthogonalityResidual5 :: Property +prop_qOrthogonalityResidual5 = forAll (genMatrix @5 @5) $ \a -> + let (q, _) = qr a + in orthogonalityResidual @5 q < 100 + +prop_qOrthogonalityResidual10 :: Property +prop_qOrthogonalityResidual10 = withMaxSuccess 20 $ forAll (genMatrix @10 @10) $ \a -> + let (q, _) = qr a + in orthogonalityResidual @10 q < 100 + +-- QR reconstruction (larger) + +prop_qrReconstruction10x10 :: Property +prop_qrReconstruction10x10 = forAll (genMatrix @10 @10) $ \a -> + let (q, r) = qr a + in normFrob (mSub a (matMul q r)) / (normFrob a + 1e-15) < 1e-6 + +prop_qrReconstruction10x5 :: Property +prop_qrReconstruction10x5 = withMaxSuccess 20 $ forAll (genMatrix @10 @5) $ \a -> + let (q, r) = qr a + in normFrob (mSub a (matMul q r)) / (normFrob a + 1e-15) < 1e-6 + +-- Givens vs Householder + +prop_givensMatchesHouseholder :: Property +prop_givensMatchesHouseholder = forAll (genMatrix @5 @5) $ \a -> + let (q1, r1) = qr a + (q2, r2) = qrGivens a + a1 = matMul q1 r1 + a2 = matMul q2 r2 + in normFrob (mSub a1 a2) / (normFrob a1 + 1e-15) < 1e-6 + +-- Least squares (property) + +prop_leastSquares10x5 :: Property +prop_leastSquares10x5 = forAll ((,) <$> genMatrix @10 @5 <*> genVector @10) $ \(a, b) -> + let x = leastSquaresQR a b + -- Compute Ax as a 10-vector + ax = matvec a x + -- Compute residual r = Ax - b + r_vec = makeVector @10 @M.P $ \i -> (ax !. i) - (b !. i) + -- Compute A^T r as a 5-vector + at = transpose a + atr = matvec at r_vec + -- Normal equations: A^T(Ax - b) should be approximately zero + in vnorm2 atr / (vnorm2 b + 1e-15) < 1e-4 diff --git a/test/Test/Residuals.hs b/test/Test/Residuals.hs new file mode 100644 --- /dev/null +++ b/test/Test/Residuals.hs @@ -0,0 +1,233 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | LAPACK-style scaled residual functions for NLA verification. +-- +-- This module provides the standard residual metrics used by LAPACK +-- (Anderson et al., 1999) and recommended by Higham (2002) for +-- verifying numerical linear algebra routines. Rather than comparing +-- element-wise against a fixed tolerance, these functions compute +-- /scaled residuals/ that account for problem size, matrix norms, +-- and machine precision. +-- +-- A scaled residual less than \(O(1)\) indicates backward stability; +-- values less than ~10--100 are considered acceptable. +-- +-- __References:__ +-- +-- * Higham, N. J. (2002). /Accuracy and Stability of Numerical +-- Algorithms/, 2nd ed., SIAM. Chapter 1. +-- * Anderson, E. et al. (1999). /LAPACK Users' Guide/, 3rd ed., SIAM. +-- * LAPACK Working Note 41: Installation Guide. +-- * Golub, G. H. & Van Loan, C. F. (2013). /Matrix Computations/, +-- 4th ed., Chapter 2. +module Test.Residuals + ( -- * Machine epsilon + machineEps + -- * Scaled residual functions + , scaledResidualLinear + , scaledResidualEigen + , scaledResidualQR + , scaledResidualSVD + , scaledResidualCholesky + , scaledResidualLU + , orthogonalityResidual + -- * Condition-number-aware tolerance + , conditionTolerance + , safeCondition2 + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix1, Ix2(..), Array) +import GHC.TypeNats (KnownNat) + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level1 (scal) +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvec) +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, transpose, mSub) +import Numeric.LinearAlgebra.Massiv.Norms (normFrob, normInf, vnorm2, vnormInf) +import Numeric.LinearAlgebra.Massiv.Eigen.SVD (singularValues) + +-- | IEEE 754 double precision machine epsilon (unit roundoff). +-- +-- \(\varepsilon = 2^{-52} \approx 2.22 \times 10^{-16}\) +-- +-- This is the smallest value such that \(1 + \varepsilon > 1\) in +-- double precision floating-point arithmetic. +machineEps :: Double +machineEps = 2.220446049250313e-16 + +-- | Scaled residual for a linear system solve \(Ax = b\). +-- +-- \[ +-- \eta(\hat{x}) = \frac{\|A\hat{x} - b\|_\infty} +-- {\|A\|_\infty \|\hat{x}\|_\infty + \|b\|_\infty} +-- \] +-- +-- Per Higham (2002), Section 7.1, this is the normwise backward error. +-- A backward-stable solver should produce \(\eta \le O(n \varepsilon)\). +scaledResidualLinear :: forall n. KnownNat n + => Matrix n n M.P Double -> Vector n M.P Double -> Vector n M.P Double -> Double +scaledResidualLinear a x b = + let ax = matvec a x + nn = dimVal @n + -- residual r = Ax - b + r = makeVector @n @M.P $ \i -> (ax !. i) - (b !. i) + numr = vnormInf r + denom = normInf a * vnormInf x + vnormInf b + fromIntegral nn * machineEps + in numr / denom + +-- | Scaled residual for an eigenpair \((\\lambda, v)\). +-- +-- \[ +-- \frac{\|Av - \lambda v\|_2}{\|A\|_F \|v\|_2} +-- \] +-- +-- Per Higham (2002), Section 14.1. A well-computed eigenpair +-- should have residual \(O(n \varepsilon)\). +scaledResidualEigen :: forall n. KnownNat n + => Matrix n n M.P Double -> Double -> Vector n M.P Double -> Double +scaledResidualEigen a lambda v = + let av = matvec a v + lambdaV = scal lambda v + r = makeVector @n @M.P $ \i -> (av !. i) - (lambdaV !. i) + numr = vnorm2 r + denom = normFrob a * vnorm2 v + machineEps + in numr / denom + +-- | Scaled residual for QR factorization \(A = QR\). +-- +-- \[ +-- \frac{\|A - QR\|_F}{\|A\|_F \cdot n \cdot \varepsilon} +-- \] +-- +-- Per LAPACK testing methodology (LAWN 41). A value less than +-- \(O(1)\) (typically < 10--100) indicates the factorization +-- is backward stable. +scaledResidualQR :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> Matrix m m M.P Double -> Matrix m n M.P Double -> Double +scaledResidualQR a q r = + let qr_prod = matMul q r + diff = mSub a qr_prod + nn = dimVal @n + numr = normFrob diff + denom = normFrob a * fromIntegral nn * machineEps + machineEps + in numr / denom + +-- | Scaled residual for SVD \(A = U \Sigma V^T\). +-- +-- \[ +-- \frac{\|A - U \Sigma V^T\|_F}{\|A\|_F \cdot \max(m,n) \cdot \varepsilon} +-- \] +-- +-- Per LAPACK testing methodology (LAWN 41). +scaledResidualSVD :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> Matrix m m M.P Double -> Vector n M.P Double + -> Matrix n n M.P Double -> Double +scaledResidualSVD a u sigma v = + let mm = dimVal @m + nn = dimVal @n + -- Build U * diag(sigma) by scaling columns of U + -- Then multiply by V^T + -- U is m x m, sigma has n entries, V is n x n + -- We need U(:,0:n-1) * diag(sigma) * V^T + -- Since our SVD returns m x m U, we take the first n columns conceptually + uSigma = makeMatrix @m @n @M.P $ \i j -> + (u ! (i, j)) * (sigma !. j) + uSigmaVt = matMul uSigma (transpose v) + diff = mSub a uSigmaVt + numr = normFrob diff + denom = normFrob a * fromIntegral (max mm nn) * machineEps + machineEps + in numr / denom + +-- | Scaled residual for Cholesky factorization \(A = GG^T\). +-- +-- \[ +-- \frac{\|A - GG^T\|_F}{\|A\|_F \cdot n \cdot \varepsilon} +-- \] +scaledResidualCholesky :: forall n. KnownNat n + => Matrix n n M.P Double -> Matrix n n M.P Double -> Double +scaledResidualCholesky a g = + let ggt = matMul g (transpose g) + diff = mSub a ggt + nn = dimVal @n + numr = normFrob diff + denom = normFrob a * fromIntegral nn * machineEps + machineEps + in numr / denom + +-- | Scaled residual for LU factorization \(PA = LU\). +-- +-- \[ +-- \frac{\|PA - LU\|_F}{\|A\|_F \cdot n \cdot \varepsilon} +-- \] +-- +-- Takes the original matrix, the packed LU matrix, and the pivot array. +-- Extracts L (unit lower triangular) and U (upper triangular) from the +-- packed form. +scaledResidualLU :: forall n. KnownNat n + => Matrix n n M.P Double -> Matrix n n M.P Double + -> Array M.P Ix1 Int -> Double +scaledResidualLU a lu_packed pivots = + let nn = dimVal @n + -- Extract L (unit lower triangular) + l = makeMatrix @n @n @M.P $ \i j -> + if i == j then 1 + else if i > j then lu_packed ! (i, j) + else 0 + -- Extract U (upper triangular) + u_mat = makeMatrix @n @n @M.P $ \i j -> + if i <= j then lu_packed ! (i, j) + else 0 + -- Construct PA by applying the permutation + pa = makeMatrix @n @n @M.P $ \i j -> + let pi_i = M.index' pivots i + in a ! (pi_i, j) + lu_prod = matMul l u_mat + diff = mSub pa lu_prod + numr = normFrob diff + denom = normFrob a * fromIntegral nn * machineEps + machineEps + in numr / denom + +-- | Orthogonality residual for a matrix \(Q\). +-- +-- \[ +-- \frac{\|Q^T Q - I\|_F}{n \cdot \varepsilon} +-- \] +-- +-- Per LAPACK testing methodology. A value less than \(O(1)\) +-- means orthogonality is preserved to within rounding error. +orthogonalityResidual :: forall n. KnownNat n + => Matrix n n M.P Double -> Double +orthogonalityResidual q = + let nn = dimVal @n + qtq = matMul (transpose q) q + diff = mSub qtq (identityMatrix @n @M.P) + numr = normFrob diff + denom = fromIntegral nn * machineEps + in numr / denom + +-- | Condition-number-aware tolerance: \(\kappa \cdot n \cdot \varepsilon\). +-- +-- For a linear system \(Ax = b\), the forward error satisfies +-- \(\|\hat{x} - x\| / \|x\| \le \kappa(A) \cdot \eta\) where +-- \(\eta\) is the backward error. A backward-stable algorithm +-- achieves \(\eta \approx n \varepsilon\), so the expected +-- forward error is \(O(\kappa \cdot n \cdot \varepsilon)\). +-- +-- See Higham (2002), Theorem 7.2. +conditionTolerance :: Double -> Int -> Double +conditionTolerance kappa n = kappa * fromIntegral n * machineEps + +-- | Estimate the 2-norm condition number \(\kappa_2(A) = \sigma_{\max} / \sigma_{\min}\). +-- +-- Uses SVD to compute singular values. Returns \(10^{16}\) for +-- numerically singular matrices (where \(\sigma_{\min} < \varepsilon \cdot \sigma_{\max}\)). +safeCondition2 :: forall n. KnownNat n => Matrix n n M.P Double -> Double +safeCondition2 a = + let sv = singularValues a + nn = dimVal @n + sigmaMax = sv !. 0 + sigmaMin = sv !. (nn - 1) + in if abs sigmaMin < machineEps * abs sigmaMax + then 1e16 + else abs sigmaMax / abs sigmaMin diff --git a/test/Test/Solve.hs b/test/Test/Solve.hs new file mode 100644 --- /dev/null +++ b/test/Test/Solve.hs @@ -0,0 +1,278 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Test.Solve (solveTests) where + +import Test.Tasty +import Test.Tasty.QuickCheck +import Test.Tasty.HUnit +import qualified Data.Massiv.Array as M + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvec) +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, transpose) +import Numeric.LinearAlgebra.Massiv.Solve.Triangular +import Numeric.LinearAlgebra.Massiv.Solve.LU +import Numeric.LinearAlgebra.Massiv.Solve.Cholesky +import Numeric.LinearAlgebra.Massiv.Solve.Banded (tridiagSolve) +import Numeric.LinearAlgebra.Massiv.Eigen.Symmetric (symmetricEigen) +import Test.Types +import Test.Residuals + +solveTests :: TestTree +solveTests = testGroup "Solve" + [ testGroup "Triangular" + [ testProperty "forward sub: Lx = b roundtrip" prop_forwardSubRoundtrip + , testProperty "back sub: Ux = b roundtrip" prop_backSubRoundtrip + ] + , testGroup "LU" + [ testProperty "PA = LU reconstruction" prop_luReconstruction + , testProperty "luSolve: Ax = b roundtrip" prop_luSolveRoundtrip + , testCase "3x3 LU known" test_luKnown + ] + , testGroup "Cholesky" + [ testProperty "A = GGᵀ reconstruction" prop_choleskyReconstruction + , testProperty "choleskySolve roundtrip" prop_choleskySolveRoundtrip + ] + , testGroup "Tridiagonal" + [ testCase "tridiag solve known" test_tridiagKnown + ] + , testGroup "Ill-conditioned" + [ testCase "LU on Hilbert 5×5" test_luHilbert + , testProperty "Cholesky near-singular SPD" prop_choleskyNearSingular + ] + , testGroup "Randsvd" + [ testProperty "luSolve with randsvd matrix" prop_luSolveRandsvd + , testProperty "choleskySolve with randsvd SPD" prop_choleskySolveRandsvd + ] + , testGroup "Scaled Residuals" + [ testProperty "LU residual bound 3×3" prop_luResidualBound + , testProperty "Cholesky residual bound 3×3" prop_choleskyResidualBound + ] + , testGroup "Cross-module" + [ testCase "det ≈ product of eigenvalues" test_detEqualsEigenProduct + ] + , testGroup "Larger Sizes" + [ testProperty "LU reconstruction 5×5" prop_luReconstruction5 + , testProperty "Cholesky reconstruction 5×5" prop_choleskyReconstruction5 + , testProperty "luSolve 10×10" prop_luSolve10 + ] + ] + +-- Triangular tests + +prop_forwardSubRoundtrip :: Property +prop_forwardSubRoundtrip = forAll ((,) <$> genLowerTriangular @3 <*> genVector @3) $ \(l, b) -> + let x = forwardSub l b + b' = matvec l x + in vecApproxEq @3 b b' + +prop_backSubRoundtrip :: Property +prop_backSubRoundtrip = forAll ((,) <$> genUpperTriangular @3 <*> genVector @3) $ \(u, b) -> + let x = backSub u b + b' = matvec u x + in vecApproxEq @3 b b' + +-- LU tests + +prop_luReconstruction :: Property +prop_luReconstruction = forAll (genMatrix @3 @3) $ \a -> + let (luMat, pivArr) = lu a + nn = 3 :: Int + l = makeMatrix @3 @3 @M.P $ \i j -> + if i == j then 1 + else if i > j then luMat ! (i, j) + else 0 :: Double + u = makeMatrix @3 @3 @M.P $ \i j -> + if i <= j then luMat ! (i, j) + else 0 :: Double + lu_ = matMul l u + -- PA + pa = makeMatrix @3 @3 @M.P $ \i j -> + a ! (M.index' pivArr i, j) + in matApproxEq @3 @3 pa lu_ + +prop_luSolveRoundtrip :: Property +prop_luSolveRoundtrip = forAll ((,) <$> genMatrix @3 @3 <*> genVector @3) $ \(a, b) -> + -- Only test for non-singular matrices (skip near-singular) + let (_, pivArr) = lu a + luMat = fst (lu a) + -- Check diagonal of U + diagOk = all (\i -> abs (luMat ! (i, i)) > 1e-6) [0..2] + in diagOk ==> + let x = luSolve a b + b' = matvec a x + in vecApproxEq @3 b b' + +test_luKnown :: Assertion +test_luKnown = do + -- A = [[2,1,1],[4,3,3],[8,7,9]] + let a = makeMatrix @3 @3 @M.P $ \i j -> case (i,j) of + (0,0) -> 2; (0,1) -> 1; (0,2) -> 1 + (1,0) -> 4; (1,1) -> 3; (1,2) -> 3 + (2,0) -> 8; (2,1) -> 7; (2,2) -> 9 + _ -> 0 :: Double + b = makeVector @3 @M.P $ \i -> case i of + 0 -> 1; 1 -> 1; _ -> 1 :: Double + x = luSolve a b + b' = matvec a x + assertBool "LU solve roundtrip" $ vecApproxEq @3 b b' + +-- Cholesky tests + +prop_choleskyReconstruction :: Property +prop_choleskyReconstruction = forAll (genSPDMatrix @3) $ \a -> + let g = cholesky a + gt = transpose g + ggt = matMul g gt + in matApproxEq @3 @3 a ggt + +prop_choleskySolveRoundtrip :: Property +prop_choleskySolveRoundtrip = forAll ((,) <$> genSPDMatrix @3 <*> genVector @3) $ \(a, b) -> + let x = choleskySolve a b + b' = matvec a x + in vecApproxEq @3 b b' + +-- Tridiagonal test + +test_tridiagKnown :: Assertion +test_tridiagKnown = do + -- Tridiagonal SPD: A = [[2,-1,0],[-1,2,-1],[0,-1,2]] + -- b = [1, 0, 1] + -- x should be [1, 1, 1] + let diag_ = makeVector @3 @M.P $ \_ -> 2 :: Double + supdiag = makeVector @3 @M.P $ \i -> if i < 2 then -1 else 0 :: Double + b = makeVector @3 @M.P $ \i -> case i of { 0 -> 1; 1 -> 0; _ -> 1 } :: Double + x = tridiagSolve diag_ supdiag b + -- Verify Ax = b manually: 2*1 + (-1)*1 = 1, (-1)*1 + 2*1 + (-1)*1 = 0, (-1)*1 + 2*1 = 1 + assertBool "x(0) ~ 1" $ (x !. 0) ~= 1 + assertBool "x(1) ~ 1" $ (x !. 1) ~= 1 + assertBool "x(2) ~ 1" $ (x !. 2) ~= 1 + +------------------------------------------------------------------------ +-- Ill-conditioned tests +------------------------------------------------------------------------ + +-- | Test 1: LU solve on the 5×5 Hilbert matrix (very ill-conditioned). +test_luHilbert :: Assertion +test_luHilbert = do + let h = hilbertMatrix @5 + x = makeVector @5 @M.P $ \i -> fromIntegral (i + 1) + b = matvec h x + x' = luSolve h b + residual = scaledResidualLinear @5 h x' b + assertBool ("LU Hilbert residual too large: " ++ show residual) + (residual < 1e-4) + +-- | Test 2: Cholesky on near-singular SPD matrix (cond = 1e4). +prop_choleskyNearSingular :: Property +prop_choleskyNearSingular = withMaxSuccess 20 $ + forAll (genSPDMatrixWithCond @5 1e4) $ \a -> + let g = cholesky a + residual = scaledResidualCholesky @5 a g + in counterexample ("Cholesky residual = " ++ show residual) $ + residual < 1000 + +------------------------------------------------------------------------ +-- Randsvd tests +------------------------------------------------------------------------ + +-- | Test 3: LU solve with a randsvd matrix (cond = 100). +prop_luSolveRandsvd :: Property +prop_luSolveRandsvd = withMaxSuccess 30 $ + forAll ((,) <$> genMatrixWithCond @5 @5 100.0 <*> genVector @5) $ \(a, b) -> + let x = luSolve a b + residual = scaledResidualLinear @5 a x b + in counterexample ("LU randsvd residual = " ++ show residual) $ + residual < 0.01 + +-- | Test 4: Cholesky solve with a randsvd SPD matrix (cond = 100). +prop_choleskySolveRandsvd :: Property +prop_choleskySolveRandsvd = withMaxSuccess 20 $ + forAll ((,) <$> genSPDMatrixWithCond @5 100.0 <*> genVector @5) $ \(a, b) -> + let x = choleskySolve a b + residual = scaledResidualLinear @5 a x b + in counterexample ("Cholesky randsvd residual = " ++ show residual) $ + residual < 0.01 + +------------------------------------------------------------------------ +-- Scaled residual bound tests +------------------------------------------------------------------------ + +-- | Test 5: LU residual bound for well-conditioned 3×3 random matrices. +prop_luResidualBound :: Property +prop_luResidualBound = forAll ((,) <$> genMatrix @3 @3 <*> genVector @3) $ \(a, b) -> + let (luMat, _) = lu a + diagOk = all (\i -> abs (luMat ! (i, i)) > 1e-6) [0..2] + in diagOk ==> + let x = luSolve a b + residual = scaledResidualLinear @3 a x b + in counterexample ("LU residual = " ++ show residual) $ + residual < 1e-6 + +-- | Test 6: Cholesky residual bound for well-conditioned 3×3 SPD matrices. +prop_choleskyResidualBound :: Property +prop_choleskyResidualBound = forAll ((,) <$> genSPDMatrix @3 <*> genVector @3) $ \(a, b) -> + let x = choleskySolve a b + residual = scaledResidualLinear @3 a x b + in counterexample ("Cholesky residual = " ++ show residual) $ + residual < 1e-6 + +------------------------------------------------------------------------ +-- Cross-module tests +------------------------------------------------------------------------ + +-- | Test 7: Determinant equals the product of eigenvalues for a known SPD matrix. +-- +-- We use a diagonal-dominant matrix: diag(4,3,2,1) + 0.1*ones(4,4). +test_detEqualsEigenProduct :: Assertion +test_detEqualsEigenProduct = do + let a = makeMatrix @4 @4 @M.P $ \i j -> + let diag_ = case i of { 0 -> 4; 1 -> 3; 2 -> 2; _ -> 1 } :: Double + in (if i == j then diag_ else 0) + 0.1 + d = det a + (eigenvals, _) = symmetricEigen a 200 1e-12 + eigenProd = product [ eigenvals !. i | i <- [0..3] ] + relErr = abs (d - eigenProd) / (abs d + 1e-15) + assertBool ("det vs eigenproduct relative error too large: " ++ show relErr) + (relErr < 1e-4) + +------------------------------------------------------------------------ +-- Larger-size tests +------------------------------------------------------------------------ + +-- | Test 8: LU reconstruction at 5×5 (like existing 3×3 test). +prop_luReconstruction5 :: Property +prop_luReconstruction5 = forAll (genMatrix @5 @5) $ \a -> + let (luMat, pivArr) = lu a + l = makeMatrix @5 @5 @M.P $ \i j -> + if i == j then 1 + else if i > j then luMat ! (i, j) + else 0 :: Double + u = makeMatrix @5 @5 @M.P $ \i j -> + if i <= j then luMat ! (i, j) + else 0 :: Double + lu_ = matMul l u + pa = makeMatrix @5 @5 @M.P $ \i j -> + a ! (M.index' pivArr i, j) + in matApproxEq @5 @5 pa lu_ + +-- | Test 9: Cholesky reconstruction at 5×5 (like existing 3×3 test). +prop_choleskyReconstruction5 :: Property +prop_choleskyReconstruction5 = forAll (genSPDMatrix @5) $ \a -> + let g = cholesky a + gt = transpose g + ggt = matMul g gt + in matApproxEq @5 @5 a ggt + +-- | Test 10: LU solve at 10×10 with diagonal guard. +prop_luSolve10 :: Property +prop_luSolve10 = withMaxSuccess 20 $ + forAll ((,) <$> genMatrix @10 @10 <*> genVector @10) $ \(a, b) -> + let (luMat, _) = lu a + diagOk = all (\i -> abs (luMat ! (i, i)) > 1e-6) [0..9] + in diagOk ==> + let x = luSolve a b + residual = scaledResidualLinear @10 a x b + in counterexample ("LU 10x10 residual = " ++ show residual) $ + residual < 1e-6 diff --git a/test/Test/Types.hs b/test/Test/Types.hs new file mode 100644 --- /dev/null +++ b/test/Test/Types.hs @@ -0,0 +1,256 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +-- | Test helpers: Arbitrary instances, approximate equality, matrix generators. +module Test.Types + ( -- * Approximate equality + (~=) + , matApproxEq + , vecApproxEq + -- * Matrix generators + , genMatrix + , genVector + , genSPDMatrix + , genUpperTriangular + , genLowerTriangular + , genSymmetric + -- * Tolerance + , defaultTol + -- * Standard test matrices + , hilbertMatrix + , wilkinsonMatrix + , frankMatrix + , hadamardMatrix + -- * Generators with controlled properties + , genMatrixWithCond + , genNearSingularMatrix + , genClusteredEigenMatrix + , genSPDMatrixWithCond + ) where + +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (Ix2(..), Sz(..), Comp(..)) +import GHC.TypeNats (KnownNat, natVal) +import Data.Proxy (Proxy(..)) +import Data.Bits (popCount) +import qualified Data.Bits as Bits +import Test.QuickCheck + +import Numeric.LinearAlgebra.Massiv.Types +import Numeric.LinearAlgebra.Massiv.Internal +import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, transpose) +import Numeric.LinearAlgebra.Massiv.Orthogonal.QR (qr) + +-- | Default tolerance for floating-point comparisons. +defaultTol :: Double +defaultTol = 1e-8 + +-- | Approximate equality for scalars. +(~=) :: Double -> Double -> Bool +x ~= y = abs (x - y) < defaultTol * (1 + abs x + abs y) + +-- | Approximate equality for matrices. +matApproxEq :: forall m n. (KnownNat m, KnownNat n) + => Matrix m n M.P Double -> Matrix m n M.P Double -> Bool +matApproxEq a b = + let r = dimVal @m + c = dimVal @n + in all (\(i, j) -> (a ! (i, j)) ~= (b ! (i, j))) + [(i, j) | i <- [0..r-1], j <- [0..c-1]] + +-- | Approximate equality for vectors. +vecApproxEq :: forall n. KnownNat n + => Vector n M.P Double -> Vector n M.P Double -> Bool +vecApproxEq a b = + let nn = dimVal @n + in all (\i -> (a !. i) ~= (b !. i)) [0..nn-1] + +-- | Generate a random m×n matrix with entries in [-10, 10]. +genMatrix :: forall m n. (KnownNat m, KnownNat n) => Gen (Matrix m n M.P Double) +genMatrix = do + let r = fromIntegral (natVal (Proxy @m)) + c = fromIntegral (natVal (Proxy @n)) + entries <- vectorOf (r * c) (choose (-10, 10)) + pure $ makeMatrix @m @n @M.P $ \i j -> entries !! (i * c + j) + +-- | Generate a random n-element vector with entries in [-10, 10]. +genVector :: forall n. KnownNat n => Gen (Vector n M.P Double) +genVector = do + let nn = fromIntegral (natVal (Proxy @n)) + entries <- vectorOf nn (choose (-10, 10)) + pure $ makeVector @n @M.P $ \i -> entries !! i + +-- | Generate a symmetric positive definite matrix: A = BBᵀ + εI. +genSPDMatrix :: forall n. KnownNat n => Gen (Matrix n n M.P Double) +genSPDMatrix = do + let nn = fromIntegral (natVal (Proxy @n)) + entries <- vectorOf (nn * nn) (choose (-5, 5)) + let b = makeMatrix @n @n @M.P $ \i j -> entries !! (i * nn + j) + -- BBᵀ + εI + epsilon = 0.1 :: Double + pure $ makeMatrix @n @n @M.P $ \i j -> + let bbT = foldl' (\acc k -> acc + (b ! (i, k)) * (b ! (j, k))) 0 [0..nn-1] + in bbT + if i == j then epsilon else 0 + +-- | Generate an upper triangular matrix with nonzero diagonal. +genUpperTriangular :: forall n. KnownNat n => Gen (Matrix n n M.P Double) +genUpperTriangular = do + let nn = fromIntegral (natVal (Proxy @n)) + entries <- vectorOf (nn * nn) (choose (-10, 10)) + diags <- vectorOf nn (choose (1, 10)) -- Ensure nonzero diagonal + pure $ makeMatrix @n @n @M.P $ \i j -> + if i == j then diags !! i + else if i < j then entries !! (i * nn + j) + else 0 + +-- | Generate a lower triangular matrix with nonzero diagonal. +genLowerTriangular :: forall n. KnownNat n => Gen (Matrix n n M.P Double) +genLowerTriangular = do + let nn = fromIntegral (natVal (Proxy @n)) + entries <- vectorOf (nn * nn) (choose (-10, 10)) + diags <- vectorOf nn (choose (1, 10)) + pure $ makeMatrix @n @n @M.P $ \i j -> + if i == j then diags !! i + else if i > j then entries !! (i * nn + j) + else 0 + +-- | Generate a symmetric matrix. +genSymmetric :: forall n. KnownNat n => Gen (Matrix n n M.P Double) +genSymmetric = do + let nn = fromIntegral (natVal (Proxy @n)) + entries <- vectorOf (nn * nn) (choose (-10, 10)) + pure $ makeMatrix @n @n @M.P $ \i j -> + if i <= j then entries !! (i * nn + j) + else entries !! (j * nn + i) + +------------------------------------------------------------------------ +-- Standard test matrices +------------------------------------------------------------------------ + +-- | Hilbert matrix of order /n/: \(H_{ij} = 1/(i+j+1)\). +-- +-- The Hilbert matrix is symmetric positive definite with an +-- exponentially growing condition number. Its inverse has integer +-- entries. +-- +-- /Reference:/ Higham (2002), Section 28.1; GVL4 p. 128. +hilbertMatrix :: forall n. KnownNat n => Matrix n n M.P Double +hilbertMatrix = makeMatrix @n @n @M.P $ \i j -> + 1 / fromIntegral (i + j + 1) + +-- | Wilkinson matrix \(W_n\): symmetric tridiagonal with +-- diagonal \(|i - \lfloor n/2 \rfloor|\) and unit sub/superdiagonal. +-- +-- Has near-degenerate eigenvalue pairs that stress eigenvalue solvers. +-- +-- /Reference:/ Higham (2002), Section 28.6; MATLAB @gallery('wilk', n)@. +wilkinsonMatrix :: forall n. KnownNat n => Matrix n n M.P Double +wilkinsonMatrix = + let nn = dimVal @n + mid = nn `div` 2 + in makeMatrix @n @n @M.P $ \i j -> + if i == j then fromIntegral (abs (i - mid)) + else if abs (i - j) == 1 then 1 + else 0 + +-- | Frank matrix of order /n/: upper Hessenberg with +-- \(\det(F) = 1\) and known positive real eigenvalues. +-- +-- \(F_{ij} = n - \max(i,j)\) for \(j \ge i-1\), zero otherwise +-- (0-indexed). +-- +-- /Reference:/ Higham (2002), Section 28.5; Frank (1958). +frankMatrix :: forall n. KnownNat n => Matrix n n M.P Double +frankMatrix = + let nn = dimVal @n + in makeMatrix @n @n @M.P $ \i j -> + if j >= i - 1 && i - 1 >= 0 || j >= i + then fromIntegral (nn - max i j) + else 0 + +-- | Hadamard matrix of order /n/ via the Sylvester\/Walsh construction. +-- +-- \(H_{ij} = (-1)^{\mathrm{popcount}(i \mathbin{\&} j)}\) +-- +-- Produces a proper Hadamard matrix when /n/ is a power of 2, +-- satisfying \(H^T H = n I\). All entries are \(\pm 1\). +-- +-- /Reference:/ Higham (2002), Section 28.3. +hadamardMatrix :: forall n. KnownNat n => Matrix n n M.P Double +hadamardMatrix = makeMatrix @n @n @M.P $ \i j -> + if even (popCount ((Bits..&.) i j)) then 1 else -1 + +------------------------------------------------------------------------ +-- Generators with controlled properties +------------------------------------------------------------------------ + +-- | Generate a matrix with prescribed 2-norm condition number /κ/. +-- +-- Uses the @randsvd@ construction: \(A = U \Sigma V^T\) where +-- \(U\) and \(V\) are random orthogonal matrices obtained from +-- QR factorization of random matrices, and the singular values +-- are geometrically spaced from 1 to \(1/\kappa\). +-- +-- /Reference:/ Higham (2002), Section 28.3; Fasi & Higham, +-- "Generating Extreme-Scale Matrices With Specified Singular +-- Values or Condition Number," SIAM J. Sci. Comput. 43(5), 2021. +genMatrixWithCond :: forall m n. (KnownNat m, KnownNat n) + => Double -> Gen (Matrix m n M.P Double) +genMatrixWithCond kappa = do + uRaw <- genMatrix @m @m + vRaw <- genMatrix @n @n + let (u, _) = qr uRaw + (v, _) = qr vRaw + mm = dimVal @m + nn = dimVal @n + minDim = min mm nn + -- Singular values geometrically spaced: sigma_0 = 1, sigma_{minDim-1} = 1/kappa + sigma = makeMatrix @m @n @M.P $ \i j -> + if i == j && i < minDim + then if minDim <= 1 then 1 + else let t = fromIntegral i / fromIntegral (minDim - 1) + in kappa ** (-t) + else 0 + pure $ matMul u (matMul sigma (transpose v)) + +-- | Generate a near-singular matrix with smallest singular value /ε/. +-- +-- Wrapper around 'genMatrixWithCond' with \(\kappa = 1/\varepsilon\). +genNearSingularMatrix :: forall n. KnownNat n + => Double -> Gen (Matrix n n M.P Double) +genNearSingularMatrix smallSV = genMatrixWithCond @n @n (1 / smallSV) + +-- | Generate a symmetric matrix with clustered eigenvalues near /c/. +-- +-- Constructs \(A = Q \Lambda Q^T\) where \(\Lambda\) is diagonal +-- with entries \(c + i \cdot 10^{-6}\) and \(Q\) is a random +-- orthogonal matrix. This stresses iterative eigenvalue solvers +-- that must resolve near-degenerate eigenvalue pairs. +genClusteredEigenMatrix :: forall n. KnownNat n + => Double -> Gen (Matrix n n M.P Double) +genClusteredEigenMatrix cluster = do + qRaw <- genMatrix @n @n + let (q, _) = qr qRaw + qt = transpose q + lambda = makeMatrix @n @n @M.P $ \i j -> + if i == j then cluster + fromIntegral i * 1e-6 else 0 + pure $ matMul q (matMul lambda qt) + +-- | Generate an SPD matrix with prescribed condition number /κ/. +-- +-- Uses the @randsvd@ construction with \(U = V\) (ensuring symmetry): +-- \(A = Q \Lambda Q^T\) where eigenvalues are geometrically spaced +-- from 1 to \(1/\kappa\). +genSPDMatrixWithCond :: forall n. KnownNat n + => Double -> Gen (Matrix n n M.P Double) +genSPDMatrixWithCond kappa = do + qRaw <- genMatrix @n @n + let (q, _) = qr qRaw + qt = transpose q + nn = dimVal @n + lambda = makeMatrix @n @n @M.P $ \i j -> + if i == j + then if nn <= 1 then 1 + else let t = fromIntegral i / fromIntegral (nn - 1) + in kappa ** (-t) + else 0 + pure $ matMul q (matMul lambda qt)