fpnla-examples-0.1: src/FPNLA/Operations/LAPACK/Strategies/POTRF/Repa/BlocksPar.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
--{-# LANGUAGE FlexibleContexts #-} --TODO sacar al terminar de debugear junto con el Show
module FPNLA.Operations.LAPACK.Strategies.POTRF.Repa.BlocksPar (
) where
import FPNLA.Matrix (Matrix,
MatrixVector,
cantCols_m,
cantRows_m,
elem_m,
fromCols_vm,
generate_m,
subMatrix_m,
toCols_vm,
transpose_m)
import FPNLA.Matrix.Instances.RepaMatrix (RepaMatrix,
RepaVector)
import FPNLA.Operations.BLAS (GEMM (gemm),
SYRK (syrk),
TRSM (trsm))
import FPNLA.Operations.LAPACK (POTRF (potrf))
import FPNLA.Operations.LAPACK.Strategies.DataTypes (CholLLVBlocksPar_Repa,
SqrBlockContext,
getSqrBlockDim)
import FPNLA.Operations.Parameters (Elt, ResM,
TransType (..),
TriangType (..),
UnitType (..),
blasResultM,
getResultDataM)
import FPNLA.Utils (iif)
import Data.Array.Repa (Array, D, Shape,
U, deepSeqArray,
delay)
import Data.Array.Repa.Eval (suspendedComputeP)
import Data.Array.Repa.Repr.Unboxed (Unbox)
import Control.DeepSeq (NFData)
--import Debug.Trace
--trace' s a = trace (s ++ ": " ++ (show a)) a
instance (Elt e, Unbox e, NFData e,
MatrixVector RepaMatrix RepaVector e,
POTRF potrfs RepaMatrix RepaVector e,
SYRK syrks RepaMatrix RepaVector e,
GEMM gemms RepaMatrix RepaVector e,
TRSM trsms RepaMatrix RepaVector e) =>
POTRF (CholLLVBlocksPar_Repa syrks gemms trsms potrfs) RepaMatrix RepaVector e where
potrf ctx (Lower mA)
= blasResultM $ chol_blk_l ctx 0 mA (generate_m 0 0 undefined)
where
chol_blk_l ctx@(block_ctx, syrk_ctx, gemm_ctx, trsm_ctx, potrf_ctx) k mA mAL
| k == 0 = let mA11' = delay . computeUnboxedP $
call_chol_unb (Lower mA11)
mA21' = delay . computeUnboxedP .
deepSeqArray mA11' .
deepSeqArray mA21 $
transpose_m $ call_trsm 1 (NoTrans . Lower $ NoUnit mA11') (transpose_m mA21)
mAx1 = delay . computeUnboxedP .
deepSeqArray mA21 $
iif (k == cantBlocks - 1) mA11' $ concatByCol_m mA11' mA21'
in chol_blk_l ctx (k + 1) mA mAx1
| k == cantBlocks = mAL
| otherwise = chol_blk_l ctx (k + 1) mA mAL'
where
{-# INLINE mA_dim #-}
mA_dim = cantCols_m mA
{-# INLINE block #-}
block = getSqrBlockDim block_ctx
{-# INLINE cantBlocks #-}
cantBlocks = mA_dim `div` block
mA10 = delay . computeUnboxedP .
deepSeqArray mAL $
subMatrix_m (block * k) 0 block (block*k) mAL
mA11 = delay . computeUnboxedP $
subMatrix_m (block * k) (block * k) block block mA
mA11' = delay . computeUnboxedP .
deepSeqArray mA10
deepSeqArray mA11 $
call_syrk (-1) (NoTrans mA10) 1 (Lower mA11)
mA11'' = delay . computeUnboxedP .
deepSeqArray mA11' $
call_chol_unb (Lower mA11')
mA20 = delay . computeUnboxedP $
subMatrix_m ((k + 1) * block) 0 (mA_dim - (k + 1)*block) (k*block) mAL
mA21 = delay . computeUnboxedP $
subMatrix_m ((k + 1) * block) (k * block) (mA_dim - (k + 1)*block) block mA
mA21' = delay . computeUnboxedP .
deepSeqArray mA20 .
deepSeqArray mA21 $
call_gemm (NoTrans mA20) (Trans mA10) (-1) 1 mA21
mA21'' = delay . computeUnboxedP .
transpose_m $ call_trsm 1 (NoTrans . Lower $ NoUnit mA11'') (transpose_m mA21')
mAx1 = delay . computeUnboxedP .
add_zeros k block . iif (k == cantBlocks - 1) mA11'' $ concatByCol_m mA11'' mA21''
mAL' = delay . computeUnboxedP .
fromCols_vm $ (toCols_vm mAL :: [RepaVector e]) ++ (toCols_vm mAx1 :: [RepaVector e])
call_syrk !m1 !alpha !beta !m2 = getResultDataM $ (syrk syrk_ctx m1 alpha beta m2 :: ResM syrks RepaVector RepaMatrix e)
call_chol_unb !m = getResultDataM $ (potrf potrf_ctx m :: ResM potrfs RepaVector RepaMatrix e)
call_gemm !m1 !m2 !alpha !beta !m3 = getResultDataM $ (gemm gemm_ctx m1 m2 alpha beta m3 :: ResM gemms RepaVector RepaMatrix e)
call_trsm !alpha !m1 !m2 = getResultDataM $ (trsm trsm_ctx alpha m1 m2 :: ResM trsms RepaVector RepaMatrix e)
add_zeros :: (Matrix RepaMatrix e) => Int -> Int -> RepaMatrix e -> RepaMatrix e
add_zeros !k !block !v = concatByCol_m (generate_m (k*block) (block) (\_ _ -> 0) :: RepaMatrix e) v
concatByCol_m :: (Unbox e, Matrix RepaMatrix e) => RepaMatrix e -> RepaMatrix e -> RepaMatrix e
concatByCol_m !m1 !m2 = let rows_m1 = cantRows_m m1
cols_m1 = cantCols_m m1
rows_m2 = cantRows_m m2
cols_m2 = cantCols_m m2
in generate_m (rows_m1 + rows_m2) cols_m1
(\i j -> iif (i >= rows_m1) (elem_m (i - rows_m1) j m2)
(elem_m i j m1))
computeUnboxedP :: (Unbox e, Shape sh) => Array D sh e -> Array U sh e
computeUnboxedP = suspendedComputeP