diff options
Diffstat (limited to 'src/Simplify.hs')
| -rw-r--r-- | src/Simplify.hs | 619 |
1 files changed, 0 insertions, 619 deletions
diff --git a/src/Simplify.hs b/src/Simplify.hs deleted file mode 100644 index 19d0c17..0000000 --- a/src/Simplify.hs +++ /dev/null @@ -1,619 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE DeriveFunctor #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ImplicitParams #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiWayIf #-} -{-# LANGUAGE QuasiQuotes #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -module Simplify ( - simplifyN, simplifyFix, - SimplifyConfig(..), defaultSimplifyConfig, simplifyWith, simplifyFixWith, -) where - -import Control.Monad (ap) -import Data.Bifunctor (first) -import Data.Function (fix) -import Data.Monoid (Any(..)) - -import Debug.Trace - -import AST -import AST.Count -import AST.Pretty -import AST.Sparse.Types -import AST.UnMonoid (acPrjCompose) -import Data -import Simplify.TH - - -data SimplifyConfig = SimplifyConfig - { scLogging :: Bool - } - -defaultSimplifyConfig :: SimplifyConfig -defaultSimplifyConfig = SimplifyConfig False - -simplifyN :: KnownEnv env => Int -> Ex env t -> Ex env t -simplifyN 0 = id -simplifyN n = simplifyN (n - 1) . simplify - -simplify :: forall env t. KnownEnv env => Ex env t -> Ex env t -simplify = - let ?accumInScope = checkAccumInScope @env knownEnv - ?config = defaultSimplifyConfig - in snd . runSM . simplify' - -simplifyWith :: forall env t. KnownEnv env => SimplifyConfig -> Ex env t -> Ex env t -simplifyWith config = - let ?accumInScope = checkAccumInScope @env knownEnv - ?config = config - in snd . runSM . simplify' - -simplifyFix :: forall env t. KnownEnv env => Ex env t -> Ex env t -simplifyFix = simplifyFixWith defaultSimplifyConfig - -simplifyFixWith :: forall env t. KnownEnv env => SimplifyConfig -> Ex env t -> Ex env t -simplifyFixWith config = - let ?accumInScope = checkAccumInScope @env knownEnv - ?config = config - in fix $ \loop e -> - let (act, e') = runSM (simplify' e) - in if act then loop e' else e' - --- | simplify monad -newtype SM tenv tt env t a = SM ((Ex env t -> Ex tenv tt) -> (Any, a)) - deriving (Functor) - -instance Applicative (SM tenv tt env t) where - pure x = SM (\_ -> (Any False, x)) - (<*>) = ap - -instance Monad (SM tenv tt env t) where - SM f >>= g = SM $ \ctx -> f ctx >>= \x -> let SM h = g x in h ctx - -runSM :: SM env t env t a -> (Bool, a) -runSM (SM f) = first getAny (f id) - -smReconstruct :: Ex env t -> SM tenv tt env t (Ex tenv tt) -smReconstruct core = SM (\ctx -> (Any False, ctx core)) - -class Monad m => ActedMonad m where - tellActed :: m () - hideActed :: m a -> m a - liftActed :: (Any, a) -> m a - -instance ActedMonad ((,) Any) where - tellActed = (Any True, ()) - hideActed (_, x) = (Any False, x) - liftActed = id - -instance ActedMonad (SM tenv tt env t) where - tellActed = SM (\_ -> tellActed) - hideActed (SM f) = SM (\ctx -> hideActed (f ctx)) - liftActed pair = SM (\_ -> pair) - --- more convenient in practice -acted :: ActedMonad m => m a -> m a -acted m = tellActed >> m - -within :: (Ex env' t' -> Ex env t) -> SM tenv tt env' t' a -> SM tenv tt env t a -within subctx (SM f) = SM $ \ctx -> f (ctx . subctx) - -simplify' :: (?accumInScope :: Bool, ?config :: SimplifyConfig, KnownEnv tenv) => Ex env t -> SM tenv tt env t (Ex env t) -simplify' expr - | scLogging ?config = do - res <- simplify'Rec expr - full <- smReconstruct res - let printed = ppExpr knownEnv full - replace a bs = concatMap (\x -> if x == a then bs else [x]) - str | '\n' `elem` printed = "--- simplify step:\n " ++ replace '\n' "\n " printed - | otherwise = "--- simplify step: " ++ printed - traceM str - return res - | otherwise = simplify'Rec expr - -simplify'Rec :: (?accumInScope :: Bool, ?config :: SimplifyConfig, KnownEnv tenv) => Ex env t -> SM tenv tt env t (Ex env t) -simplify'Rec = \case - -- inlining - ELet _ rhs body - | cheapExpr rhs - -> acted $ simplify' (substInline rhs body) - - | Occ lexOcc runOcc <- occCount IZ body - , ((not ?accumInScope || not (hasAdds rhs)) && lexOcc <= One && runOcc <= One) -- without effects, normal rules apply - || (lexOcc == One && runOcc == One) -- with effects, linear inlining is still allowed, but weakening is not - -> acted $ simplify' (substInline rhs body) - - -- let splitting / let peeling - ELet _ (EPair _ a b) body -> - acted $ simplify' $ - ELet ext a $ - ELet ext (weakenExpr WSink b) $ - subst (\_ t -> \case IZ -> EPair ext (EVar ext (typeOf a) (IS IZ)) (EVar ext (typeOf b) IZ) - IS i -> EVar ext t (IS (IS i))) - body - ELet _ (EJust _ a) body -> - acted $ simplify' $ ELet ext a $ subst0 (EJust ext (EVar ext (typeOf a) IZ)) body - ELet _ (EInl _ t2 a) body -> - acted $ simplify' $ ELet ext a $ subst0 (EInl ext t2 (EVar ext (typeOf a) IZ)) body - ELet _ (EInr _ t1 a) body -> - acted $ simplify' $ ELet ext a $ subst0 (EInr ext t1 (EVar ext (typeOf a) IZ)) body - - -- let rotation - ELet _ (ELet _ rhs a) b -> do - b' <- within (ELet ext (ELet ext rhs a)) $ simplify' b - acted $ simplify' $ - ELet ext rhs $ - ELet ext a $ - weakenExpr (WCopy WSink) b' - - -- beta rules for products - EFst _ (EPair _ e e') - | not (hasAdds e') -> acted $ simplify' e - | otherwise -> acted $ simplify' $ ELet ext e' (weakenExpr WSink e) - ESnd _ (EPair _ e' e) - | not (hasAdds e') -> acted $ simplify' e - | otherwise -> acted $ simplify' $ ELet ext e' (weakenExpr WSink e) - - -- beta rules for coproducts - ECase _ (EInl _ _ e) rhs _ -> acted $ simplify' (ELet ext e rhs) - ECase _ (EInr _ _ e) _ rhs -> acted $ simplify' (ELet ext e rhs) - - -- beta rules for maybe - EMaybe _ e1 _ ENothing{} -> acted $ simplify' e1 - EMaybe _ _ e1 (EJust _ e2) -> acted $ simplify' $ ELet ext e2 e1 - - -- let floating - EFst _ (ELet _ rhs body) -> acted $ simplify' (ELet ext rhs (EFst ext body)) - ESnd _ (ELet _ rhs body) -> acted $ simplify' (ELet ext rhs (ESnd ext body)) - ECase _ (ELet _ rhs body) e1 e2 -> acted $ simplify' (ELet ext rhs (ECase ext body (weakenExpr (WCopy WSink) e1) (weakenExpr (WCopy WSink) e2))) - EIdx0 _ (ELet _ rhs body) -> acted $ simplify' (ELet ext rhs (EIdx0 ext body)) - EIdx1 _ (ELet _ rhs body) e -> acted $ simplify' (ELet ext rhs (EIdx1 ext body (weakenExpr WSink e))) - EAccum _ t p e1 sp (ELet _ rhs body) acc -> - acted $ simplify' $ - ELet ext rhs $ - EAccum ext t p (weakenExpr WSink e1) sp body (weakenExpr WSink acc) - - -- let () = e in () ~> e - ELet _ e1 (ENil _) | STNil <- typeOf e1 -> - acted $ simplify' e1 - - -- map (\_ -> x) e ~> build (shape e) (\_ -> x) - EMap _ e1 e2 - | Occ Zero Zero <- occCount IZ e1 - , STArr n _ <- typeOf e2 -> - acted $ simplify' $ - EBuild ext n (EShape ext e2) $ - subst (\_ t' -> \case IZ -> error "Unused variable was used" - IS i -> EVar ext t' (IS i)) - e1 - - -- vertical fusion - EMap _ e1 (EMap _ e2 e3) -> - acted $ simplify' $ - EMap ext (ELet ext e2 (weakenExpr (WCopy WSink) e1)) e3 - - -- projection down-commuting - EFst _ (ECase _ e1 e2 e3) -> - acted $ simplify' $ - ECase ext e1 (EFst ext e2) (EFst ext e3) - ESnd _ (ECase _ e1 e2 e3) -> - acted $ simplify' $ - ECase ext e1 (ESnd ext e2) (ESnd ext e3) - EFst _ (EMaybe _ e1 e2 e3) -> - acted $ simplify' $ - EMaybe ext (EFst ext e1) (EFst ext e2) e3 - ESnd _ (EMaybe _ e1 e2 e3) -> - acted $ simplify' $ - EMaybe ext (ESnd ext e1) (ESnd ext e2) e3 - - -- TODO: more array indexing - EIdx _ (EBuild _ _ e1 e2) e3 | not (hasAdds e1), not (hasAdds e2) -> acted $ simplify' $ elet e3 e2 - EIdx _ (EMap _ e1 e2) e3 | not (hasAdds e1) -> acted $ simplify' $ elet (EIdx ext e2 e3) e1 - EIdx _ (EReplicate1Inner _ e1 e2) e3 | not (hasAdds e1) -> acted $ simplify' $ EIdx ext e2 (EFst ext e3) - EIdx _ (EUnit _ e1) e2 | not (hasAdds e2) -> acted $ simplify' $ e1 - - -- TODO: more array shape - EShape _ (EBuild _ _ e1 e2) | not (hasAdds e2) -> acted $ simplify' e1 - EShape _ (EMap _ e1 e2) | not (hasAdds e1) -> acted $ simplify' (EShape ext e2) - - -- TODO: more constant folding - EOp _ OIf (EConst _ STBool True) -> acted $ return (EInl ext STNil (ENil ext)) - EOp _ OIf (EConst _ STBool False) -> acted $ return (EInr ext STNil (ENil ext)) - - -- inline cheap array constructors - ELet _ (EReplicate1Inner _ e1 e2) e3 -> - acted $ simplify' $ - ELet ext (EPair ext e1 e2) $ - let v = EVar ext (STPair tIx (typeOf e2)) IZ - in subst0 (EReplicate1Inner ext (EFst ext v) (ESnd ext v)) e3 - -- -- TODO: This is a bad idea and anyway only helps in practice if (!) is - -- -- cheap, which it can't be because (!) is not cheap if you do AD after. - -- -- Should do proper SoA representation. - -- ELet _ (EBuild _ n e1 e2) e3 | cheapExpr e2 -> - -- acted $ simplify' $ - -- ELet ext e1 $ - -- subst0 (EBuild ext n (EVar ext (tTup (sreplicate n tIx)) IZ) (weakenExpr (WCopy WSink) e2)) e3 - - -- eta rule for unit - e | STNil <- typeOf e, not ?accumInScope || not (hasAdds e) -> - case e of - ENil _ -> return e - _ -> acted $ return (ENil ext) - - EBuild _ SZ _ e -> - acted $ simplify' $ EUnit ext (substInline (ENil ext) e) - - -- monoid rules - EAccum _ t p e1 sp e2 acc -> do - e1' <- within (\e1' -> EAccum ext t p e1' sp e2 acc ) $ simplify' e1 - e2' <- within (\e2' -> EAccum ext t p e1' sp e2' acc ) $ simplify' e2 - acc' <- within (\acc' -> EAccum ext t p e1' sp e2' acc') $ simplify' acc - simplifyOHT (OneHotTerm SAID t p e1' sp e2') - (acted $ return (ENil ext)) - (\sp' (InContext w wrap e) -> do - e' <- within (\e' -> wrap $ EAccum ext t SAPHere (ENil ext) sp' e' (weakenExpr w acc')) $ simplify' e - return (wrap $ EAccum ext t SAPHere (ENil ext) sp' e' (weakenExpr w acc'))) - (\(InContext w wrap (OneHotTerm _ t' p' e1'' sp' e2'')) -> do - -- The acted management here is a hideous mess. - e1''' <- hideActed $ within (\e1''' -> wrap $ EAccum ext t' p' e1''' sp' e2'' (weakenExpr w acc')) $ simplify' e1'' - e2''' <- hideActed $ within (\e2''' -> wrap $ EAccum ext t' p' e1''' sp' e2''' (weakenExpr w acc')) $ simplify' e2'' - return (wrap $ EAccum ext t' p' e1''' sp' e2''' (weakenExpr w acc'))) - EPlus _ _ (EZero _ _ _) e -> acted $ simplify' e - EPlus _ _ e (EZero _ _ _) -> acted $ simplify' e - EOneHot _ t p e1 e2 -> do - e1' <- within (\e1' -> EOneHot ext t p e1' e2 ) $ simplify' e1 - e2' <- within (\e2' -> EOneHot ext t p e1' e2') $ simplify' e2 - simplifyOHT (OneHotTerm SAIS t p e1' (spDense (acPrjTy p t)) e2') - (acted $ return (EZero ext t (zeroInfoFromOneHot t p e1 e2))) - (\sp' (InContext _ wrap e) -> - case isDense t sp' of - Just Refl -> do - e' <- hideActed $ within wrap $ simplify' e - return (wrap e') - Nothing -> error "simplifyOneHotTerm sparsified a dense Sparse") - (\(InContext _ wrap (OneHotTerm _ t' p' e1'' sp' e2'')) -> - case isDense (acPrjTy p' t') sp' of - Just Refl -> do - e1''' <- hideActed $ within (\e1''' -> wrap $ EOneHot ext t' p' e1''' e2'') $ simplify' e1'' - e2''' <- hideActed $ within (\e2''' -> wrap $ EOneHot ext t' p' e1''' e2''') $ simplify' e2'' - return (wrap $ EOneHot ext t' p' e1''' e2''') - Nothing -> error "simplifyOneHotTerm sparsified a dense Sparse") - - -- type-specific equations for plus - EPlus _ SMTNil e1 e2 | not (hasAdds e1), not (hasAdds e2) -> - acted $ return (ENil ext) - - EPlus _ (SMTPair t1 t2) (EPair _ a1 b1) (EPair _ a2 b2) -> - acted $ simplify' $ EPair ext (EPlus ext t1 a1 a2) (EPlus ext t2 b1 b2) - - EPlus _ (SMTLEither t1 _) (ELInl _ dt2 a1) (ELInl _ _ a2) -> - acted $ simplify' $ ELInl ext dt2 (EPlus ext t1 a1 a2) - EPlus _ (SMTLEither _ t2) (ELInr _ dt1 b1) (ELInr _ _ b2) -> - acted $ simplify' $ ELInr ext dt1 (EPlus ext t2 b1 b2) - EPlus _ SMTLEither{} ELNil{} e -> acted $ simplify' e - EPlus _ SMTLEither{} e ELNil{} -> acted $ simplify' e - - EPlus _ (SMTMaybe t) (EJust _ e1) (EJust _ e2) -> - acted $ simplify' $ EJust ext (EPlus ext t e1 e2) - EPlus _ SMTMaybe{} ENothing{} e -> acted $ simplify' e - EPlus _ SMTMaybe{} e ENothing{} -> acted $ simplify' e - - -- fallback recursion - EVar _ t i -> pure $ EVar ext t i - ELet _ a b -> [simprec| ELet ext *a *b |] - EPair _ a b -> [simprec| EPair ext *a *b |] - EFst _ e -> [simprec| EFst ext *e |] - ESnd _ e -> [simprec| ESnd ext *e |] - ENil _ -> pure $ ENil ext - EInl _ t e -> [simprec| EInl ext t *e |] - EInr _ t e -> [simprec| EInr ext t *e |] - ECase _ e a b -> [simprec| ECase ext *e *a *b |] - ENothing _ t -> pure $ ENothing ext t - EJust _ e -> [simprec| EJust ext *e |] - EMaybe _ a b e -> [simprec| EMaybe ext *a *b *e |] - ELNil _ t1 t2 -> pure $ ELNil ext t1 t2 - ELInl _ t e -> [simprec| ELInl ext t *e |] - ELInr _ t e -> [simprec| ELInr ext t *e |] - ELCase _ e a b c -> [simprec| ELCase ext *e *a *b *c |] - EConstArr _ n t v -> pure $ EConstArr ext n t v - EBuild _ n a b -> [simprec| EBuild ext n *a *b |] - EMap _ a b -> [simprec| EMap ext *a *b |] - EFold1Inner _ cm a b c -> [simprec| EFold1Inner ext cm *a *b *c |] - ESum1Inner _ e -> [simprec| ESum1Inner ext *e |] - EUnit _ e -> [simprec| EUnit ext *e |] - EReplicate1Inner _ a b -> [simprec| EReplicate1Inner ext *a *b |] - EMaximum1Inner _ e -> [simprec| EMaximum1Inner ext *e |] - EMinimum1Inner _ e -> [simprec| EMinimum1Inner ext *e |] - EReshape _ n a b -> [simprec| EReshape ext n *a *b |] - EZip _ a b -> [simprec| EZip ext *a *b |] - EFold1InnerD1 _ cm a b c -> [simprec| EFold1InnerD1 ext cm *a *b *c |] - EFold1InnerD2 _ cm a b c -> [simprec| EFold1InnerD2 ext cm *a *b *c |] - EConst _ t v -> pure $ EConst ext t v - EIdx0 _ e -> [simprec| EIdx0 ext *e |] - EIdx1 _ a b -> [simprec| EIdx1 ext *a *b |] - EIdx _ a b -> [simprec| EIdx ext *a *b |] - EShape _ e -> [simprec| EShape ext *e |] - EOp _ op e -> [simprec| EOp ext op *e |] - ECustom _ s t p a b c e1 e2 -> do - a' <- within (\a' -> ECustom ext s t p a' b c e1 e2) (let ?accumInScope = False in simplify' a) - b' <- within (\b' -> ECustom ext s t p a' b' c e1 e2) (let ?accumInScope = False in simplify' b) - c' <- within (\c' -> ECustom ext s t p a' b' c' e1 e2) (let ?accumInScope = False in simplify' c) - e1' <- within (\e1' -> ECustom ext s t p a' b' c' e1' e2) (simplify' e1) - e2' <- within (\e2' -> ECustom ext s t p a' b' c' e1' e2') (simplify' e2) - pure (ECustom ext s t p a' b' c' e1' e2') - ERecompute _ e -> [simprec| ERecompute ext *e |] - EWith _ t e1 e2 -> do - e1' <- within (\e1' -> EWith ext t e1' e2) (simplify' e1) - e2' <- within (\e2' -> EWith ext t e1' e2') (let ?accumInScope = True in simplify' e2) - pure (EWith ext t e1' e2') - -- EOneHot _ t p e1 e2 -> [simprec| EOneHot ext t p *e1 *e2 |] - -- EAccum _ t p e1 sp e2 acc -> [simprec| EAccum ext t p *e1 sp *e2 *acc |] - EZero _ t e -> [simprec| EZero ext t *e |] - EDeepZero _ t e -> [simprec| EDeepZero ext t *e |] - EPlus _ t a b -> [simprec| EPlus ext t *a *b |] - EError _ t s -> pure $ EError ext t s - --- | This can be made more precise by tracking (and not counting) adds on --- locally eliminated accumulators. -hasAdds :: Expr x env t -> Bool -hasAdds = \case - EVar _ _ _ -> False - ELet _ rhs body -> hasAdds rhs || hasAdds body - EPair _ a b -> hasAdds a || hasAdds b - EFst _ e -> hasAdds e - ESnd _ e -> hasAdds e - ENil _ -> False - EInl _ _ e -> hasAdds e - EInr _ _ e -> hasAdds e - ECase _ e a b -> hasAdds e || hasAdds a || hasAdds b - ENothing _ _ -> False - EJust _ e -> hasAdds e - EMaybe _ a b e -> hasAdds a || hasAdds b || hasAdds e - ELNil _ _ _ -> False - ELInl _ _ e -> hasAdds e - ELInr _ _ e -> hasAdds e - ELCase _ e a b c -> hasAdds e || hasAdds a || hasAdds b || hasAdds c - EConstArr _ _ _ _ -> False - EBuild _ _ a b -> hasAdds a || hasAdds b - EMap _ a b -> hasAdds a || hasAdds b - EFold1Inner _ _ a b c -> hasAdds a || hasAdds b || hasAdds c - ESum1Inner _ e -> hasAdds e - EUnit _ e -> hasAdds e - EReplicate1Inner _ a b -> hasAdds a || hasAdds b - EMaximum1Inner _ e -> hasAdds e - EMinimum1Inner _ e -> hasAdds e - EReshape _ _ a b -> hasAdds a || hasAdds b - EZip _ a b -> hasAdds a || hasAdds b - EFold1InnerD1 _ _ a b c -> hasAdds a || hasAdds b || hasAdds c - EFold1InnerD2 _ _ a b c -> hasAdds a || hasAdds b || hasAdds c - ECustom _ _ _ _ a b c d e -> hasAdds a || hasAdds b || hasAdds c || hasAdds d || hasAdds e - EConst _ _ _ -> False - EIdx0 _ e -> hasAdds e - EIdx1 _ a b -> hasAdds a || hasAdds b - EIdx _ a b -> hasAdds a || hasAdds b - EShape _ e -> hasAdds e - EOp _ _ e -> hasAdds e - EWith _ _ a b -> hasAdds a || hasAdds b - ERecompute _ e -> hasAdds e - EAccum _ _ _ _ _ _ _ -> True - EZero _ _ e -> hasAdds e - EDeepZero _ _ e -> hasAdds e - EPlus _ _ a b -> hasAdds a || hasAdds b - EOneHot _ _ _ a b -> hasAdds a || hasAdds b - EError _ _ _ -> False - -checkAccumInScope :: SList STy env -> Bool -checkAccumInScope = \case SNil -> False - SCons t env -> check t || checkAccumInScope env - where - check :: STy t -> Bool - check STNil = False - check (STPair s t) = check s || check t - check (STEither s t) = check s || check t - check (STLEither s t) = check s || check t - check (STMaybe t) = check t - check (STArr _ t) = check t - check (STScal _) = False - check STAccum{} = True - -data OneHotTerm dense env a where - OneHotTerm :: SAIDense dense -> SMTy a -> SAcPrj p a b -> Ex env (AcIdx dense p a) -> Sparse b c -> Ex env c -> OneHotTerm dense env a -deriving instance Show (OneHotTerm dense env a) - -data InContext f env (a :: Ty) where - InContext :: env :> env' -> (forall t. Ex env' t -> Ex env t) -> f env' a -> InContext f env a - -simplifyOHT_recogniseMonoid :: ActedMonad m => OneHotTerm dense env a -> m (OneHotTerm dense env a) -simplifyOHT_recogniseMonoid (OneHotTerm dense t prj idx sp val) = do - val' <- liftActed $ recogniseMonoid (applySparse sp (acPrjTy prj t)) val - return $ OneHotTerm dense t prj idx sp val' - -simplifyOHT_unsparse :: ActedMonad m => OneHotTerm dense env a -> m (InContext (OneHotTerm dense) env a) -simplifyOHT_unsparse (OneHotTerm SAID t prj1 idx1 sp1 val1) = - unsparseOneHotD sp1 val1 $ \w wrap prj2 idx2 sp2 val2 -> - acPrjCompose SAID prj1 (weakenExpr w idx1) prj2 idx2 $ \prj' idx' -> - return $ InContext w wrap (OneHotTerm SAID t prj' idx' sp2 val2) -simplifyOHT_unsparse oht@(OneHotTerm SAIS _ _ _ _ _) = return $ InContext WId id oht - -simplifyOHT_concat :: ActedMonad m => OneHotTerm dense env a -> m (OneHotTerm dense env a) -simplifyOHT_concat (OneHotTerm @dense @_ @_ @_ @env dense t1 prj1 idx1 sp (EOneHot @_ @c @p2 _ t2 prj2 idx2 val)) - | Just Refl <- isDense (acPrjTy prj1 t1) sp = - let idx2' :: Ex env (AcIdx dense p2 c) - idx2' = case dense of - SAID -> reduceAcIdx t2 prj2 idx2 - SAIS -> idx2 - in acPrjCompose dense prj1 idx1 prj2 idx2' $ \prj' idx' -> - acted $ return $ OneHotTerm dense t1 prj' idx' (spDense (acPrjTy prj' t1)) val -simplifyOHT_concat oht = return oht - --- -- Property not expressed in types: if the Sparse in the input OneHotTerm is --- -- dense, then the Sparse in the output will also be dense. This property is --- -- used when simplifying EOneHot, which cannot represent sparsity. -simplifyOHT :: ActedMonad m => OneHotTerm dense env a - -> m r -- ^ Zero case (onehot is actually zero) - -> (forall b. Sparse a b -> InContext Ex env b -> m r) -- ^ Trivial case (no zeros in onehot) - -> (InContext (OneHotTerm dense) env a -> m r) -- ^ Simplified - -> m r -simplifyOHT oht kzero ktriv k = do - -- traceM $ "sOHT: input " ++ show oht - oht1 <- simplifyOHT_recogniseMonoid oht - -- traceM $ "sOHT: recog " ++ show oht1 - InContext w1 wrap1 oht2 <- simplifyOHT_unsparse oht1 - -- traceM $ "sOHT: unspa " ++ show oht2 - oht3 <- simplifyOHT_concat oht2 - -- traceM $ "sOHT: conca " ++ show oht3 - -- traceM "" - case oht3 of - OneHotTerm _ _ _ _ _ EZero{} -> kzero - OneHotTerm _ _ SAPHere _ sp val -> ktriv sp (InContext w1 wrap1 val) - _ -> k (InContext w1 wrap1 oht3) - --- Sets the acted flag whenever a non-trivial projection is returned or the --- output Sparse is different from the input Sparse. -unsparseOneHotD :: ActedMonad m => Sparse a a' -> Ex env a' - -> (forall p b c env'. env :> env' -> (forall s. Ex env' s -> Ex env s) - -> SAcPrj p a b -> Ex env' (AcIdxD p a) -> Sparse b c -> Ex env' c -> m r) -> m r -unsparseOneHotD topsp topval k = case (topsp, topval) of - -- eliminate always-Just sparse onehot - (SpSparse s, EOneHot _ (SMTMaybe t) (SAPJust prj) idx val) -> - acted $ unsparseOneHotD s (EOneHot ext t prj idx val) k - - -- expand the top levels of a onehot for a sparse type into a onehot for the - -- corresponding non-sparse type - (SpPair s1 _, EOneHot _ (SMTPair t1 _) (SAPFst prj) idx val) -> - unsparseOneHotD s1 (EOneHot ext t1 prj (efst idx) val) $ \w wrap spprj idx' s1' e' -> - acted $ k w wrap (SAPFst spprj) idx' s1' e' - (SpPair _ s2, EOneHot _ (SMTPair _ t2) (SAPSnd prj) idx val) -> - unsparseOneHotD s2 (EOneHot ext t2 prj (esnd idx) val) $ \w wrap spprj idx' s1' e' -> - acted $ k w wrap (SAPSnd spprj) idx' s1' e' - (SpLEither s1 _, EOneHot _ (SMTLEither t1 _) (SAPLeft prj) idx val) -> - unsparseOneHotD s1 (EOneHot ext t1 prj idx val) $ \w wrap spprj idx' s1' e' -> - acted $ k w wrap (SAPLeft spprj) idx' s1' e' - (SpLEither _ s2, EOneHot _ (SMTLEither _ t2) (SAPRight prj) idx val) -> - unsparseOneHotD s2 (EOneHot ext t2 prj idx val) $ \w wrap spprj idx' s1' e' -> - acted $ k w wrap (SAPRight spprj) idx' s1' e' - (SpMaybe s1, EOneHot _ (SMTMaybe t1) (SAPJust prj) idx val) -> - unsparseOneHotD s1 (EOneHot ext t1 prj idx val) $ \w wrap spprj idx' s1' e' -> - acted $ k w wrap (SAPJust spprj) idx' s1' e' - (SpArr s1, EOneHot _ (SMTArr _ t1) (SAPArrIdx prj) idx val) - | Dict <- styKnown (typeOf idx) -> - unsparseOneHotD s1 (EOneHot ext t1 prj (esnd (evar IZ)) (weakenExpr WSink val)) $ \w wrap spprj idx' s1' e' -> - acted $ k (w .> WSink) (elet idx . wrap) (SAPArrIdx spprj) (EPair ext (efst (efst (evar (w @> IZ)))) idx') s1' e' - - -- anything else we don't know how to improve - _ -> k WId id SAPHere (ENil ext) topsp topval - -{- -unsparseOneHotS :: ActedMonad m - => Sparse a a' -> Ex env a' - -> (forall b. Sparse a b -> Ex env b -> m r) -> m r -unsparseOneHotS topsp topval k = case (topsp, topval) of - -- order is relevant to make sure we set the acted flag correctly - (SpAbsent, v@ENil{}) -> k SpAbsent v - (SpAbsent, v@EZero{}) -> k SpAbsent v - (SpAbsent, _) -> acted $ k SpAbsent (EZero ext SMTNil (ENil ext)) - (_, EZero{}) -> acted $ k SpAbsent (EZero ext SMTNil (ENil ext)) - (sp, _) | isAbsent sp -> acted $ k SpAbsent (EZero ext SMTNil (ENil ext)) - - -- the unsparsifying - (SpSparse s, EOneHot _ (SMTMaybe t) (SAPJust prj) idx val) -> - acted $ unsparseOneHotS s (EOneHot ext t prj idx val) k - - -- recursion - -- TODO: coproducts could safely become projections as they do not need - -- zeroinfo. But that would only work if the coproduct is at the top, because - -- as soon as we hit a product, we need zeroinfo to make it a projection and - -- we don't have that. - (SpSparse s, e) -> k (SpSparse s) e - (SpPair s1 _, EOneHot _ (SMTPair t1 _) (SAPFst prj) idx val) -> - unsparseOneHotS s1 (EOneHot ext t1 prj (efst idx) val) $ \s1' e' -> - acted $ k (SpPair s1' SpAbsent) (EPair ext e' (ENil ext)) - (SpPair _ s2, EOneHot _ (SMTPair _ t2) (SAPSnd prj) idx val) -> - unsparseOneHotS s2 (EOneHot ext t2 prj (esnd idx) val) $ \s2' e' -> - acted $ k (SpPair SpAbsent s2') (EPair ext (ENil ext) e') - (SpLEither s1 s2, EOneHot _ (SMTLEither t1 _) (SAPLeft prj) idx val) -> - unsparseOneHotS s1 (EOneHot ext t1 prj idx val) $ \s1' e' -> do - case s2 of SpAbsent -> pure () ; _ -> tellActed - k (SpLEither s1' SpAbsent) (ELInl ext STNil e') - (SpLEither s1 s2, EOneHot _ (SMTLEither _ t2) (SAPRight prj) idx val) -> - unsparseOneHotS s2 (EOneHot ext t2 prj idx val) $ \s2' e' -> do - case s1 of SpAbsent -> pure () ; _ -> tellActed - acted $ k (SpLEither SpAbsent s2') (ELInr ext STNil e') - (SpMaybe s1, EOneHot _ (SMTMaybe t1) (SAPJust prj) idx val) -> - unsparseOneHotS s1 (EOneHot ext t1 prj idx val) $ \s1' e' -> - k (SpMaybe s1') (EJust ext e') - (SpArr s1, EOneHot _ (SMTArr n t1) (SAPArrIdx prj) idx val) -> - unsparseOneHotS s1 (EOneHot ext t1 prj (esnd (evar IZ)) (weakenExpr WSink val)) $ \s1' e' -> - k (SpArr s1') (elet idx $ EOneHot ext (SMTArr n (applySparse s1' _)) (SAPArrIdx SAPHere) (EPair ext (efst (evar IZ)) (ENil ext)) e') - _ -> _ --} - --- | Recognises 'EZero' and 'EOneHot'. -recogniseMonoid :: SMTy t -> Ex env t -> (Any, Ex env t) -recogniseMonoid _ e@EOneHot{} = return e -recogniseMonoid SMTNil (ENil _) = acted $ return $ EZero ext SMTNil (ENil ext) -recogniseMonoid typ@(SMTPair t1 t2) (EPair _ a b) = - ((,) <$> recogniseMonoid t1 a <*> recogniseMonoid t2 b) >>= \case - (EZero _ _ ezi1, EZero _ _ ezi2) -> acted $ return $ EZero ext typ (EPair ext ezi1 ezi2) - (a', EZero _ _ ezi2) -> acted $ EOneHot ext typ (SAPFst SAPHere) (EPair ext (ENil ext) ezi2) <$> recogniseMonoid t1 a' - (EZero _ _ ezi1, b') -> acted $ EOneHot ext typ (SAPSnd SAPHere) (EPair ext ezi1 (ENil ext)) <$> recogniseMonoid t2 b' - (a', b') -> return $ EPair ext a' b' -recogniseMonoid typ@(SMTLEither t1 t2) expr = - case expr of - ELNil{} -> acted $ return $ EZero ext typ (ENil ext) - ELInl _ _ e -> acted $ EOneHot ext typ (SAPLeft SAPHere) (ENil ext) <$> recogniseMonoid t1 e - ELInr _ _ e -> acted $ EOneHot ext typ (SAPRight SAPHere) (ENil ext) <$> recogniseMonoid t2 e - _ -> return expr -recogniseMonoid typ@(SMTMaybe t1) expr = - case expr of - ENothing{} -> acted $ return $ EZero ext typ (ENil ext) - EJust _ e -> acted $ EOneHot ext typ (SAPJust SAPHere) (ENil ext) <$> recogniseMonoid t1 e - _ -> return expr -recogniseMonoid typ@(SMTArr SZ t) (EUnit _ e) = - acted $ do - e' <- recogniseMonoid t e - return $ - ELet ext e' $ - EOneHot ext typ (SAPArrIdx SAPHere) - (EPair ext (EPair ext (ENil ext) (EUnit ext (makeZeroInfo t (EVar ext (fromSMTy t) IZ)))) - (ENil ext)) - (EVar ext (fromSMTy t) IZ) -recogniseMonoid typ@(SMTScal sty) e@(EConst _ _ x) = case (sty, x) of - (STI32, 0) -> acted $ return $ EZero ext typ (ENil ext) - (STI64, 0) -> acted $ return $ EZero ext typ (ENil ext) - (STF32, 0) -> acted $ return $ EZero ext typ (ENil ext) - (STF64, 0) -> acted $ return $ EZero ext typ (ENil ext) - _ -> return e -recogniseMonoid _ e = return e - -reduceAcIdx :: SMTy a -> SAcPrj p a b -> Ex env (AcIdxS p a) -> Ex env (AcIdxD p a) -reduceAcIdx topty topprj e = case (topty, topprj) of - (_, SAPHere) -> ENil ext - (SMTPair t1 _, SAPFst p) -> reduceAcIdx t1 p (efst e) - (SMTPair _ t2, SAPSnd p) -> reduceAcIdx t2 p (esnd e) - (SMTLEither t1 _ , SAPLeft p) -> reduceAcIdx t1 p e - (SMTLEither _ t2, SAPRight p) -> reduceAcIdx t2 p e - (SMTMaybe t1, SAPJust p) -> reduceAcIdx t1 p e - (SMTArr _ t, SAPArrIdx p) -> - eunPair e $ \_ e1 e2 -> - EPair ext (efst e1) (reduceAcIdx t p e2) - -zeroInfoFromOneHot :: SMTy t -> SAcPrj p t a -> Ex env (AcIdxS p t) -> Ex env a -> Ex env (ZeroInfo t) -zeroInfoFromOneHot = \ty prj eidx e -> ELet ext eidx $ go ty prj (EVar ext (typeOf eidx) IZ) (weakenExpr WSink e) - where - -- invariant: AcIdx expression is duplicable - go :: SMTy t -> SAcPrj p t a -> Ex env (AcIdxS p t) -> Ex env a -> Ex env (ZeroInfo t) - go t SAPHere _ e = makeZeroInfo t e - go (SMTPair t1 _) (SAPFst prj) eidx e = EPair ext (go t1 prj (EFst ext eidx) e) (ESnd ext eidx) - go (SMTPair _ t2) (SAPSnd prj) eidx e = EPair ext (EFst ext eidx) (go t2 prj (ESnd ext eidx) e) - go SMTLEither{} _ _ _ = ENil ext - go SMTMaybe{} _ _ _ = ENil ext - go SMTArr{} SAPArrIdx{} eidx _ = ESnd ext (EFst ext eidx) |