ADPfusion-0.6.0.0: ADP/Fusion/Core/SynVar/Fill.hs
module ADP.Fusion.Core.SynVar.Fill where
import Control.Monad
import Control.Monad.Morph (hoist, MFunctor (..))
import Control.Monad.Primitive
import Control.Monad.ST
import Control.Monad.Trans.Class (lift, MonadTrans (..))
import Control.Monad (when,forM_)
import Data.Dynamic
import Data.List (nub,sort,group)
import Data.Maybe (fromJust)
import Data.Proxy
import Data.Type.Equality
import Data.Vector.Fusion.Util (Id(..))
import Debug.Trace (traceShow)
import GHC.Exts (inline)
import GHC.TypeNats
import qualified Data.Data as D
import qualified Data.List as L
import qualified Data.Typeable as T
import qualified Data.Vector as V
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified GHC.Generics as G
import System.IO.Unsafe
import System.CPUTime
import GHC.Conc (pseq)
import Data.PrimitiveArray
import ADP.Fusion.Core.SynVar.Array -- TODO we want to keep only classes in here, move instances to the corresponding modules
import ADP.Fusion.Core.SynVar.Recursive.Type
import ADP.Fusion.Core.SynVar.TableWrap
import Debug.Trace
{-
-- | A vanilla context-free grammar
data CFG
-- | This grammar is a multi-cfg in a monotone setting
data MonotoneMCFG
-- * Unsafely mutate 'ITbls' and similar tables in the forward phase.
-- | Mutate a cell in a stack of syntactic variables.
--
-- TODO generalize to monad morphism via @mmorph@ package. This will allow
-- more interesting @mrph@ functions that can, for example, track some
-- state in the forward phase. (Note that this can be dangerous, we do
-- /not/ want to have this state influence forward results, unless that can
-- be made deterministic, or we'll break Bellman)
class MutateCell (h ∷ *) (s ∷ *) (im ∷ * → *) i where
mutateCell
∷ (Monad om, PrimMonad om)
⇒ Proxy h
→ Int
→ Int
→ (forall a . im a → om a)
→ s
→ LimitType i
→ i
→ om ()
-- |
class MutateTables (h :: *) (s :: *) (im :: * -> *) where
mutateTables :: (Monad om, PrimMonad om) => Proxy h -> (forall a . im a -> om a) -> s -> om s
class TableOrder (s :: *) where
tableLittleOrder :: s -> [Int]
tableBigOrder :: s -> [Int]
instance TableOrder Z where
tableLittleOrder Z = []
tableBigOrder Z = []
{-# Inline tableLittleOrder #-}
{-# Inline tableBigOrder #-}
instance
( TableOrder ts
, KnownNat bo
-- , KnownNat lo
) ⇒ TableOrder (ts:.TwITbl bo im arr c i x) where
tableLittleOrder (ts:.TW (ITbl tlo _ _) _) =
let -- tlo = fromIntegral $ natVal (Proxy ∷ Proxy lo)
in tlo : tableLittleOrder ts
tableBigOrder (ts:.TW (ITbl _ _ _) _) =
let tbo = fromIntegral $ natVal (Proxy ∷ Proxy bo)
in tbo : tableBigOrder ts
{-# Inline tableLittleOrder #-}
{-# Inline tableBigOrder #-}
-- | @IRec@s do not need an order, given that they do not memoize.
instance (TableOrder ts) => TableOrder (ts:.TwIRec im c i x) where
tableLittleOrder (ts:._) = tableLittleOrder ts
tableBigOrder (ts:._) = tableBigOrder ts
{-# Inline tableLittleOrder #-}
{-# Inline tableBigOrder #-}
-- ** individual instances for filling a *single cell*
instance
(
) => MutateCell p Z im i where
mutateCell _ _ _ _ Z _ _ = return ()
{-# INLINE mutateCell #-}
instance
( MutateCell CFG ts im i
) => MutateCell CFG (ts:.TwIRec im c i x) im i where
mutateCell h bo lo mrph (ts:._) lu i = do
mutateCell h bo lo mrph ts lu i
{-# Inline mutateCell #-}
instance
( PrimArrayOps arr i x
, MPrimArrayOps arr i x
, MutateCell CFG ts im i
, KnownNat bo
-- , KnownNat lo
) => MutateCell CFG (ts:.TwITbl bo im arr c i x) im i where
mutateCell h bo lo mrph (ts:.TW (ITbl tlo c arr) f) lu i = do
let tbo = fromIntegral $ natVal (Proxy ∷ Proxy bo)
-- tlo = fromIntegral $ natVal (Proxy ∷ Proxy lo)
mutateCell h bo lo mrph ts lu i
when (bo==tbo && lo==tlo) $ do
marr <- unsafeThaw arr
z <- (inline mrph) $ f lu i
writeM marr i z
{-# INLINE mutateCell #-}
{-
- TODOThe following code goes into ADPfusionSubword!
-
type ZS2 = Z:.Subword I:.Subword I
instance
( PrimArrayOps arr ZS2 x
, MPrimArrayOps arr ZS2 x
, MutateCell MonotoneMCFG ts im ZS2
) => MutateCell MonotoneMCFG (ts:.TwITbl im arr c ZS2 x) im ZS2 where
mutateCell h bo lo mrph (ts:.TW (ITbl tbo tlo c arr) f) lu iklj@(Z:.Subword (i:.k):.Subword(l:.j)) = do
mutateCell h bo lo mrph ts lu iklj
when (bo==tbo && lo==tlo && k<=l) $ do
marr <- unsafeThaw arr
z <- (inline mrph) $ f lu iklj
writeM marr iklj z
{-# INLINE mutateCell #-}
instance
( PrimArrayOps arr (Subword I) x
, MPrimArrayOps arr (Subword I) x
, MutateCell h ts im (Z:.Subword I:.Subword I)
) => MutateCell h (ts:.TwITbl im arr c (Subword I) x) im (Z:.Subword I:.Subword I) where
mutateCell h bo lo mrph (ts:.TW (ITbl tbo tlo c arr) f) lu@(Z:.Subword (l:._):.Subword(_:.u)) ix@(Z:.Subword (i1:.j1):.Subword (i2:.j2)) = do
mutateCell h bo lo mrph ts lu ix
when (bo==tbo && lo==tlo && i1==i2 && j1==j2) $ do
let i = i1
let j = j1
marr <- unsafeThaw arr
z <- (inline mrph) $ f (subword l u) (subword i j)
writeM marr (subword i j) z
{-# Inline mutateCell #-}
-}
-- ** individual instances for filling a complete table and extracting the
-- bounds
instance
( MutateCell h (ts:.TwITbl bo im arr c i x) im i
, PrimArrayOps arr i x
, Show i
, IndexStream i
, TableOrder (ts:.TwITbl bo im arr c i x)
) => MutateTables h (ts:.TwITbl bo im arr c i x) im where
mutateTables h mrph tt@(_:.TW (ITbl lo _ arr) _) = do
let to = upperBound arr
-- TODO (1) find the set of orders for the synvars
let !tbos = VU.fromList . nub . sort $ tableBigOrder tt
let !tlos = VU.fromList . nub . sort $ tableLittleOrder tt
VU.forM_ tbos $ \bo ->
case (VU.length tlos) of
1 -> let lo = VU.head tlos
in flip SM.mapM_ (streamUp zeroBound' to) $ \k ->
mutateCell h bo lo (inline mrph) tt to k
-- TODO each big-order group should be allowed to have its own sets
-- of bounds. within a group, it doesn't make a lot of sense to
-- have different bounds? Is there a use case for that even?
_ -> flip SM.mapM_ (streamUp zeroBound' to) $ \k ->
VU.forM_ tlos $ \lo ->
mutateCell h bo lo (inline mrph) tt to k
return tt
{-# INLINE mutateTables #-}
-- | Default table filling, assuming that the forward monad is just @IO@.
--
-- TODO generalize to @MonadIO@ or @MonadPrim@.
mutateTablesDefault :: MutateTables CFG t Id => t -> t
mutateTablesDefault t = unsafePerformIO $ mutateTables (Proxy :: Proxy CFG) (return . unId) t
{-# INLINE mutateTablesDefault #-}
-- | Mutate tables, but observe certain hints. We use this for monotone
-- mcfgs for now.
mutateTablesWithHints :: MutateTables h t Id => Proxy h -> t -> t
mutateTablesWithHints h t = unsafePerformIO $ mutateTables h (return . unId) t
mutateTablesST t = runST $ mutateTablesNew t
{-# Inline mutateTablesST #-}
class CountNumberOfCells t where
countNumberOfCells ∷ t → Integer
instance CountNumberOfCells Z where
countNumberOfCells Z = 0
instance
( CountNumberOfCells ts
, Index i
, PrimArrayOps arr i x
) ⇒ CountNumberOfCells (ts:.TwITbl bo Id arr c i x) where
countNumberOfCells (ts:.(TW (ITbl lo _ arr) fun)) =
countNumberOfCells ts + (product . totalSize $ upperBound arr)
data PerfCounter = PerfCounter
{ picoSeconds :: !Integer
, seconds :: !Double
, numberOfCells :: !Integer
}
deriving (Eq,Ord,Show)
data Mutated ts = Mutated
{ mutatedTables ∷ !ts
, perfCounter ∷ !PerfCounter
, eachBigPerfCounter ∷ [PerfCounter]
}
-- |
--
-- TODO new way how to do table filling. Because we now have heterogeneous
-- tables (i) group tables by @big order@ into different bins; (ii) check
-- that each bin has the same bounds (needed? -- could we have
-- smaller-sized tables once in a while); (iii) run each bin one after the
-- other
--
-- TODO measure performance penalty, if any. We might need liberal
-- INLINEABLE, and specialization. On the other hand, we can do the
-- freeze/unfreeze outside of table filling.
mutateTablesNew
:: forall t m .
( TableOrder t
, TSBO t
, Monad m
, PrimMonad m
, CountNumberOfCells t
)
=> t
-> m (Mutated t)
mutateTablesNew ts = do
-- sort the tables according to [bigorder,type,littleorder]. For each
-- @bigorder@, we should have only one @type@ and can therefor do the
-- following (i) get subset of the @ts@, (ii) use outermost of @ts@ to
-- get bounds, (iii) fill these tables
-- let !tbos = VU.fromList . nub . sort $ tableBigOrder ts
let justOrder = L.map (\d → (qBigOrder d, qLittleOrder d))
let ds = L.sort $ asDyn ts
let goM ∷ (Monad m, PrimMonad m) ⇒ [Q] → [PerfCounter] → m [PerfCounter]
goM [] ps = return $ reverse ps
goM xs ps = do
(ys,p) <- fillWithDyn xs ts
goM ys (p:ps)
{-# Inlinable goM #-}
startTime ← unsafeIOToPrim getCPUTime
ps ← goM ds []
stopTime ← unsafeIOToPrim getCPUTime
let deltaTime = max 1 $ stopTime - startTime
return $! Mutated
{ mutatedTables = ts
, perfCounter = PerfCounter
{ picoSeconds = deltaTime
, seconds = 1e-12 * fromIntegral deltaTime
, numberOfCells = countNumberOfCells ts
}
, eachBigPerfCounter = ps
}
{-# Inline mutateTablesNew #-}
data Q = Q
{ qBigOrder :: Int
, qLittleOrder :: Int
, qTypeRep :: T.TypeRep
, qObject :: Dynamic
, qTable :: Dynamic
, qFunction :: Dynamic
}
deriving (Show)
instance Eq Q where
Q bo1 lo1 tr1 _ _ _ == Q bo2 lo2 tr2 _ _ _ = (bo1,tr1,lo1) == (bo2,tr2,lo2)
instance Ord Q where
Q bo1 lo1 tr1 _ _ _ `compare` Q bo2 lo2 tr2 _ _ _ = (bo1,lo1,tr1) `compare` (bo2,lo2,tr2)
-- | Find the outermost table that has a certain big order and then fill
-- from there.
class TSBO t where
asDyn :: t -> [Q]
fillWithDyn :: (Monad m, PrimMonad m) => [Q] -> t -> m ([Q], PerfCounter)
instance TSBO Z where
asDyn Z = []
fillWithDyn qs Z = return (qs, PerfCounter 0 0 0)
{-# Inlinable asDyn #-}
{-# Inline fillWithDyn #-}
instance
( TSBO ts
, Typeable arr
, Typeable c
, Typeable i
, Typeable x
, PrimArrayOps arr i x
, MPrimArrayOps arr i x
, IndexStream i
, KnownNat bo
-- , KnownNat lo
) => TSBO (ts:.TwITbl bo Id arr c i x) where
asDyn (ts:.t@(TW (ITbl lo _ arr) fun)) =
let bo = fromIntegral $ natVal (Proxy ∷ Proxy bo)
-- lo = fromIntegral $ natVal (Proxy ∷ Proxy lo)
in Q bo lo (T.typeOf t) (toDyn t) (toDyn arr) (seq fun $ toDyn fun) : asDyn ts
fillWithDyn qs (ts:.t@(TW (ITbl _ _ arrDirect) fDirect)) = do
let to = upperBound arrDirect
bo = fromIntegral $ natVal (Proxy ∷ Proxy bo)
-- lo = fromIntegral $ natVal (Proxy ∷ Proxy lo)
-- @hs@ are all tables that can be filled here
-- @ns@ are all tables we can't fill and need to process further down
-- the line
-- TODO FIXME FIXME FIXME why are the typereps different???
let (hs,ns) = L.span (\Q{..} -> qBigOrder == bo) qs -- && qTypeRep == T.typeOf t) qs
if null hs
then fillWithDyn qs ts
else do
let ms = Prelude.map concreteTW hs
af = Prelude.map concreteAF hs
concreteTW = (maybe (error "fromDynamic should not fail!")
(\x -> x `asTypeOf` t)
. fromDynamic . qObject)
concreteAF q = ( (`asTypeOf` arrDirect) . fromJust . fromDynamic $ qTable q
, (`asTypeOf` fDirect) . fromJust . fromDynamic $ qFunction q
)
-- We have a single table and should short-circuit here
--
-- TODO we should specialize for tables of lengh @1..k@ for some
-- small k. For @1@ and Needleman-Wunsch, we have a very nice @1.8@
-- seconds down to @1.25@ seconds. :-)
--
-- TODO how about
-- case ms of
-- [a] -> bla
-- [a,b] -> bla
-- [a,b,c] -> bla
-- [a:b:c:d:ms'] -> bla >> go ms'
-- measure if this yields meaningful performance improvements
--
-- TODO also consider if we maybe just put marrfs into a vector
--
-- TODO we should use TH here.
--
-- (1) Have @Proxy @0@, say to set up big and small orders -- this
-- gives us the order on the type level. @data One = One, data Two
-- = Two, ...@ might be easier... maybe this is not too annoying to
-- write using type equality
--
-- (2) Then deconstruct the @ts:.t@ things with TH into the correct
-- pieces.
--
-- (3) Finally generate fill code. This should yield to performance
-- similar to what we have here with the @case of 1@ construction,
-- because @fDirect@ is partially floated out.
--
marrfs <- V.fromList <$> Prelude.mapM (\(TW (ITbl _ _ arr) f) -> unsafeThaw arr >>= \marr -> return (marr,f)) ms
startTime ← unsafeIOToPrim getCPUTime
case (V.length marrfs) of
1 -> do -- let (!marr,!f) = marrfs V.! 0 -- this takes 1.3 seconds for NeedlemanWunsch
-- marr <- unsafeThaw arrDirect -- this takes 0.8 seconds for NeedlemanWunsch
marr <- unsafeThaw arrDirect -- (fst $ af!!0) -- this takes 1.3 seconds for NeedlemanWunsch
let !ffff = fDirect --snd $ af!!0
flip SM.mapM_ (streamUp zeroBound' to) $ \k -> do
-- TODO @inline mrph@ ...
z <- (return . unId) $ fDirect to k
writeM marr k z
-- We have more than one table in will work over the list of tables
_ -> do flip SM.mapM_ (streamUp zeroBound' to) $ \k ->
V.forM_ marrfs $ \(marr,f) -> do
z <- (return . unId) $ f to k
writeM marr k z
-- traceShow (hs,length ms) $
stopTime ← unsafeIOToPrim getCPUTime
let deltaTime = stopTime - startTime
let perf = PerfCounter
{ picoSeconds = deltaTime
, seconds = 1e-12 * fromIntegral deltaTime
, numberOfCells = sum $ Prelude.map (\(TW t _) → product . totalSize . upperBound $ iTblArray t) ms
}
return (ns, perf)
{-# Inline fillWithDyn #-}
-- We don't need to capture @IRec@ tables as no table-filling takes place
-- for those tables. @asDyn@ therefore just collects on the remaining @ts@,
-- while @fillWithDyn@ hands of to the next possible table.
instance
( TSBO ts
) => TSBO (ts:.TwIRec Id c i x) where
asDyn (ts:.t@(TW (IRec _ _) _)) = asDyn ts
fillWithDyn qs (ts:._) = fillWithDyn qs ts
{-# Inlinable asDyn #-}
{-# Inline fillWithDyn #-}
-}