linear-massiv-0.1.0.0: src/Numeric/LinearAlgebra/Massiv/Internal/Kernel.hs
{-# 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 #-}