Merge branch 'sparse'HEADmaster

1 files changed, 108 insertions, 5 deletions

diff --git a/src/AST/UnMonoid.hs b/src/AST/UnMonoid.hs
index f5841e0..48dd709 100644
--- a/src/AST/UnMonoid.hs
+++ b/src/AST/UnMonoid.hs
@@ -1,18 +1,22 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE TypeOperators #-}
-module AST.UnMonoid (unMonoid, zero, plus) where
+module AST.UnMonoid (unMonoid, zero, plus, acPrjCompose) where
 
 import AST
+import AST.Sparse.Types
 import Data
 
 
--- | Remove 'EZero', 'EPlus' and 'EOneHot' from the program by expanding them
--- into their concrete implementations.
+-- | Remove 'EZero', 'EDeepZero', 'EPlus' and 'EOneHot' from the program by
+-- expanding them into their concrete implementations. Also ensure that
+-- 'EAccum' has a dense sparsity.
 unMonoid :: Ex env t -> Ex env t
 unMonoid = \case
   EZero _ t e -> zero t e
+  EDeepZero _ t e -> deepZero t e
   EPlus _ t a b -> plus t (unMonoid a) (unMonoid b)
   EOneHot _ t p a b -> onehot t p (unMonoid a) (unMonoid b)
 
@@ -49,7 +53,10 @@ unMonoid = \case
   ECustom _ t1 t2 t3 a b c e1 e2 -> ECustom ext t1 t2 t3 (unMonoid a) (unMonoid b) (unMonoid c) (unMonoid e1) (unMonoid e2)
   ERecompute _ e -> ERecompute ext (unMonoid e)
   EWith _ t a b -> EWith ext t (unMonoid a) (unMonoid b)
-  EAccum _ t p a b e -> EAccum ext t p (unMonoid a) (unMonoid b) (unMonoid e)
+  EAccum _ t p eidx sp eval eacc ->
+    accumulateSparse (acPrjTy p t) sp eval $ \w prj2 idx2 val2 ->
+    acPrjCompose SAID p (weakenExpr w eidx) prj2 idx2 $ \prj' idx' ->
+      EAccum ext t prj' (unMonoid idx') (spDense (acPrjTy prj' t)) (unMonoid val2) (weakenExpr w (unMonoid eacc))
   EError _ t s -> EError ext t s
 
 zero :: SMTy t -> Ex env (ZeroInfo t) -> Ex env t
@@ -67,6 +74,27 @@ zero (SMTScal t) _ = case t of
   STF32 -> EConst ext STF32 0.0
   STF64 -> EConst ext STF64 0.0
 
+deepZero :: SMTy t -> Ex env (DeepZeroInfo t) -> Ex env t
+deepZero SMTNil e = elet e $ ENil ext
+deepZero (SMTPair t1 t2) e =
+  ELet ext e $ EPair ext (deepZero t1 (EFst ext (EVar ext (typeOf e) IZ)))
+                         (deepZero t2 (ESnd ext (EVar ext (typeOf e) IZ)))
+deepZero (SMTLEither t1 t2) e =
+  elcase e
+    (ELNil ext (fromSMTy t1) (fromSMTy t2))
+    (ELInl ext (fromSMTy t2) (deepZero t1 (evar IZ)))
+    (ELInr ext (fromSMTy t1) (deepZero t2 (evar IZ)))
+deepZero (SMTMaybe t) e =
+  emaybe e
+    (ENothing ext (fromSMTy t))
+    (EJust ext (deepZero t (evar IZ)))
+deepZero (SMTArr _ t) e = emap (deepZero t (evar IZ)) e
+deepZero (SMTScal t) _ = case t of
+  STI32 -> EConst ext STI32 0
+  STI64 -> EConst ext STI64 0
+  STF32 -> EConst ext STF32 0.0
+  STF64 -> EConst ext STF64 0.0
+
 plus :: SMTy t -> Ex env t -> Ex env t -> Ex env t
 -- don't destroy the effects!
 plus SMTNil a b = ELet ext a $ ELet ext (weakenExpr WSink b) $ ENil ext
@@ -107,7 +135,7 @@ plus (SMTArr _ t) a b =
            a b
 plus (SMTScal t) a b = EOp ext (OAdd t) (EPair ext a b)
 
-onehot :: SMTy t -> SAcPrj p t a -> Ex env (AcIdx p t) -> Ex env a -> Ex env t
+onehot :: SMTy t -> SAcPrj p t a -> Ex env (AcIdxS p t) -> Ex env a -> Ex env t
 onehot typ topprj idx arg = case (typ, topprj) of
   (_, SAPHere) ->
     ELet ext arg $
@@ -145,3 +173,78 @@ onehot typ topprj idx arg = case (typ, topprj) of
              (onehot t1 prj (ESnd ext (EVar ext (typeOf idx) (IS IZ))) (weakenExpr (WSink .> WSink) arg))
              (ELet ext (EIdx ext (ESnd ext (EFst ext (EVar ext (typeOf idx) (IS IZ)))) (EVar ext tidx IZ)) $
                 zero t1 (EVar ext (tZeroInfo t1) IZ))
+
+accumulateSparse
+  :: SMTy t -> Sparse t t' -> Ex env t'
+  -> (forall p b env'. env :> env' -> SAcPrj p t b -> Ex env' (AcIdxD p t) -> Ex env' b -> Ex env' TNil)
+  -> Ex env TNil
+accumulateSparse topty topsp arg accum = case (topty, topsp) of
+  (_, s) | Just Refl <- isDense topty s ->
+    accum WId SAPHere (ENil ext) arg
+  (SMTScal _, SpScal) ->
+    accum WId SAPHere (ENil ext) arg  -- should be handled by isDense already, but meh
+  (_, SpSparse s) ->
+    emaybe arg
+      (ENil ext)
+      (accumulateSparse topty s (evar IZ) (\w -> accum (WPop w)))
+  (_, SpAbsent) ->
+    ENil ext
+  (SMTPair t1 t2, SpPair s1 s2) ->
+    eunPair arg $ \w1 e1 e2 ->
+      elet (accumulateSparse t1 s1 e1 (\w prj -> accum (w .> w1) (SAPFst prj))) $
+        accumulateSparse t2 s2 (weakenExpr WSink e2) (\w prj -> accum (w .> WSink .> w1) (SAPSnd prj))
+  (SMTLEither t1 t2, SpLEither s1 s2) ->
+    elcase arg
+      (ENil ext)
+      (accumulateSparse t1 s1 (evar IZ) (\w prj -> accum (WPop w) (SAPLeft prj)))
+      (accumulateSparse t2 s2 (evar IZ) (\w prj -> accum (WPop w) (SAPRight prj)))
+  (SMTMaybe t, SpMaybe s) ->
+    emaybe arg
+      (ENil ext)
+      (accumulateSparse t s (evar IZ) (\w prj -> accum (WPop w) (SAPJust prj)))
+  (SMTArr n t, SpArr s) ->
+    let tn = tTup (sreplicate n tIx) in
+    elet arg $
+    elet (EBuild ext n (EShape ext (evar IZ)) $
+            accumulateSparse t s
+              (EIdx ext (evar (IS IZ)) (EVar ext tn IZ))
+              (\w prj idx val -> accum (WPop (WPop w)) (SAPArrIdx prj) (EPair ext (EVar ext tn (w @> IZ)) idx) val)) $
+      ENil ext
+
+acPrjCompose
+  :: SAIDense dense
+  -> SAcPrj p1 a b -> Ex env (AcIdx dense p1 a)
+  -> SAcPrj p2 b c -> Ex env (AcIdx dense p2 b)
+  -> (forall p'. SAcPrj p' a c -> Ex env (AcIdx dense p' a) -> r) -> r
+acPrjCompose _ SAPHere _ p2 idx2 k = k p2 idx2
+acPrjCompose SAID (SAPFst p1) idx1 p2 idx2 k =
+  acPrjCompose SAID p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPFst p') idx'
+acPrjCompose SAID (SAPSnd p1) idx1 p2 idx2 k =
+  acPrjCompose SAID p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPSnd p') idx'
+acPrjCompose SAIS (SAPFst p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAIS p1 (efst (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPFst p') (elet idx1 $ EPair ext idx' (esnd (evar IZ)))
+acPrjCompose SAIS (SAPSnd p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAIS p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPSnd p') (elet idx1 $ EPair ext (efst (evar IZ)) idx')
+acPrjCompose d (SAPLeft p1) idx1 p2 idx2 k =
+  acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPLeft p') idx'
+acPrjCompose d (SAPRight p1) idx1 p2 idx2 k =
+  acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPRight p') idx'
+acPrjCompose d (SAPJust p1) idx1 p2 idx2 k =
+  acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPJust p') idx'
+acPrjCompose SAID (SAPArrIdx p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAID p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPArrIdx p') (elet idx1 $ EPair ext (EFst ext (EVar ext (typeOf idx1) IZ)) idx')
+acPrjCompose SAIS (SAPArrIdx p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAIS p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPArrIdx p') (elet idx1 $ EPair ext (EFst ext (EVar ext (typeOf idx1) IZ)) idx')