-- | This module implements facilities for determining whether a
-- reduction or fold can be expressed in a closed form (i.e. not as a
-- SOAC).
--
-- Right now, the module can detect only trivial cases.  In the
-- future, we would like to make it more powerful, as well as possibly
-- also being able to analyse sequential loops.
module Futhark.Optimise.Simplify.Rules.ClosedForm
  ( foldClosedForm,
    loopClosedForm,
  )
where

import Control.Monad
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Construct
import Futhark.IR
import Futhark.Optimise.Simplify.Rule
import Futhark.Optimise.Simplify.Rules.Simple (VarLookup)
import Futhark.Transform.Rename

{-
Motivation:

  let {*[int,x_size_27] map_computed_shape_1286} = replicate(x_size_27,
                                                             all_equal_shape_1044) in
  let {*[bool,x_size_27] map_size_checks_1292} = replicate(x_size_27, x_1291) in
  let {bool all_equal_checked_1298, int all_equal_shape_1299} =
    reduceT(fn {bool, int} (bool bacc_1293, int nacc_1294, bool belm_1295,
                            int nelm_1296) =>
              let {bool tuplit_elems_1297} = bacc_1293 && belm_1295 in
              {tuplit_elems_1297, nelm_1296},
            {True, 0}, map_size_checks_1292, map_computed_shape_1286)
-}

-- | @foldClosedForm look foldfun accargs arrargs@ determines whether
-- each of the results of @foldfun@ can be expressed in a closed form.
foldClosedForm ::
  (BuilderOps rep) =>
  VarLookup rep ->
  Pat (LetDec rep) ->
  Lambda rep ->
  [SubExp] ->
  [VName] ->
  RuleM rep ()
foldClosedForm :: forall rep.
BuilderOps rep =>
VarLookup rep
-> Pat (LetDec rep)
-> Lambda rep
-> [SubExp]
-> [VName]
-> RuleM rep ()
foldClosedForm VarLookup rep
look Pat (LetDec rep)
pat Lambda rep
lam [SubExp]
accs [VName]
arrs = do
  SubExp
inputsize <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> RuleM rep [TypeBase Shape NoUniqueness] -> RuleM rep SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> RuleM rep (TypeBase Shape NoUniqueness))
-> [VName] -> RuleM rep [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> RuleM rep (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs

  PrimType
t <- case Pat (LetDec rep) -> [TypeBase Shape NoUniqueness]
forall dec. Typed dec => Pat dec -> [TypeBase Shape NoUniqueness]
patTypes Pat (LetDec rep)
pat of
    [Prim FloatType {}] -> RuleM rep PrimType
forall rep a. RuleM rep a
cannotSimplify
    [Prim PrimType
t] -> PrimType -> RuleM rep PrimType
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
t
    [TypeBase Shape NoUniqueness]
_ -> RuleM rep PrimType
forall rep a. RuleM rep a
cannotSimplify

  Body rep
closedBody <-
    [VName]
-> SubExp
-> Names
-> IntType
-> Map VName SubExp
-> [VName]
-> Body rep
-> [SubExp]
-> RuleM rep (Body rep)
forall rep.
BuilderOps rep =>
[VName]
-> SubExp
-> Names
-> IntType
-> Map VName SubExp
-> [VName]
-> Body rep
-> [SubExp]
-> RuleM rep (Body rep)
checkResults
      (Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat)
      SubExp
inputsize
      Names
forall a. Monoid a => a
mempty
      IntType
Int64
      Map VName SubExp
knownBnds
      ((Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName (Lambda rep -> [Param (LParamInfo rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam))
      (Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
lam)
      [SubExp]
accs
  VName
isEmpty <- String -> RuleM rep VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"fold_input_is_empty"
  [VName] -> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
isEmpty] (Exp (Rep (RuleM rep)) -> RuleM rep ())
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> Exp (Rep (RuleM rep))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (RuleM rep)))
-> BasicOp -> Exp (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$
      CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
inputsize (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0)
  Pat (LetDec (Rep (RuleM rep)))
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec rep)
Pat (LetDec (Rep (RuleM rep)))
pat
    (Exp rep -> RuleM rep ()) -> RuleM rep (Exp rep) -> RuleM rep ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( [SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
forall rep.
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match [VName -> SubExp
Var VName
isEmpty]
            ([Case (Body rep)]
 -> Body rep -> MatchDec (BranchType rep) -> Exp rep)
-> RuleM rep [Case (Body rep)]
-> RuleM rep (Body rep -> MatchDec (BranchType rep) -> Exp rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Case (Body rep) -> [Case (Body rep)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Case (Body rep) -> [Case (Body rep)])
-> (Body rep -> Case (Body rep)) -> Body rep -> [Case (Body rep)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe PrimValue] -> Body rep -> Case (Body rep)
forall body. [Maybe PrimValue] -> body -> Case body
Case [PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
True] (Body rep -> [Case (Body rep)])
-> RuleM rep (Body rep) -> RuleM rep [Case (Body rep)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExp] -> RuleM rep (Body (Rep (RuleM rep)))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp]
accs)
            RuleM rep (Body rep -> MatchDec (BranchType rep) -> Exp rep)
-> RuleM rep (Body rep)
-> RuleM rep (MatchDec (BranchType rep) -> Exp rep)
forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Body rep -> RuleM rep (Body rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody Body rep
closedBody
            RuleM rep (MatchDec (BranchType rep) -> Exp rep)
-> RuleM rep (MatchDec (BranchType rep)) -> RuleM rep (Exp rep)
forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MatchDec (BranchType rep) -> RuleM rep (MatchDec (BranchType rep))
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([BranchType rep] -> MatchSort -> MatchDec (BranchType rep)
forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec [PrimType -> BranchType rep
forall rt. IsBodyType rt => PrimType -> rt
primBodyType PrimType
t] MatchSort
MatchNormal)
        )
  where
    knownBnds :: Map VName SubExp
knownBnds = VarLookup rep
-> Lambda rep -> [SubExp] -> [VName] -> Map VName SubExp
forall rep.
VarLookup rep
-> Lambda rep -> [SubExp] -> [VName] -> Map VName SubExp
determineKnownBindings VarLookup rep
look Lambda rep
lam [SubExp]
accs [VName]
arrs

-- | @loopClosedForm pat respat merge bound bodys@ determines whether
-- the do-loop can be expressed in a closed form.
loopClosedForm ::
  (BuilderOps rep) =>
  Pat (LetDec rep) ->
  [(FParam rep, SubExp)] ->
  Names ->
  IntType ->
  SubExp ->
  Body rep ->
  RuleM rep ()
loopClosedForm :: forall rep.
BuilderOps rep =>
Pat (LetDec rep)
-> [(FParam rep, SubExp)]
-> Names
-> IntType
-> SubExp
-> Body rep
-> RuleM rep ()
loopClosedForm Pat (LetDec rep)
pat [(FParam rep, SubExp)]
merge Names
i IntType
it SubExp
bound Body rep
body = do
  PrimType
t <- case Pat (LetDec rep) -> [TypeBase Shape NoUniqueness]
forall dec. Typed dec => Pat dec -> [TypeBase Shape NoUniqueness]
patTypes Pat (LetDec rep)
pat of
    [Prim FloatType {}] -> RuleM rep PrimType
forall rep a. RuleM rep a
cannotSimplify
    [Prim PrimType
t] -> PrimType -> RuleM rep PrimType
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
t
    [TypeBase Shape NoUniqueness]
_ -> RuleM rep PrimType
forall rep a. RuleM rep a
cannotSimplify

  Body rep
closedBody <-
    [VName]
-> SubExp
-> Names
-> IntType
-> Map VName SubExp
-> [VName]
-> Body rep
-> [SubExp]
-> RuleM rep (Body rep)
forall rep.
BuilderOps rep =>
[VName]
-> SubExp
-> Names
-> IntType
-> Map VName SubExp
-> [VName]
-> Body rep
-> [SubExp]
-> RuleM rep (Body rep)
checkResults
      [VName]
mergenames
      SubExp
bound
      Names
i
      IntType
it
      Map VName SubExp
knownBnds
      ((Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
mergeidents)
      Body rep
body
      [SubExp]
mergeexp
  VName
isEmpty <- String -> RuleM rep VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"bound_is_zero"
  [VName] -> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
isEmpty] (Exp (Rep (RuleM rep)) -> RuleM rep ())
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> Exp (Rep (RuleM rep))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (RuleM rep)))
-> BasicOp -> Exp (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$
      CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSlt IntType
it) SubExp
bound (IntType -> Integer -> SubExp
intConst IntType
it Integer
0)

  Pat (LetDec (Rep (RuleM rep)))
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec rep)
Pat (LetDec (Rep (RuleM rep)))
pat
    (Exp rep -> RuleM rep ()) -> RuleM rep (Exp rep) -> RuleM rep ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( [SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
forall rep.
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match [VName -> SubExp
Var VName
isEmpty]
            ([Case (Body rep)]
 -> Body rep -> MatchDec (BranchType rep) -> Exp rep)
-> RuleM rep [Case (Body rep)]
-> RuleM rep (Body rep -> MatchDec (BranchType rep) -> Exp rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Case (Body rep) -> [Case (Body rep)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Case (Body rep) -> [Case (Body rep)])
-> (Body rep -> Case (Body rep)) -> Body rep -> [Case (Body rep)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe PrimValue] -> Body rep -> Case (Body rep)
forall body. [Maybe PrimValue] -> body -> Case body
Case [PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (Bool -> PrimValue
BoolValue Bool
True)] (Body rep -> [Case (Body rep)])
-> RuleM rep (Body rep) -> RuleM rep [Case (Body rep)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExp] -> RuleM rep (Body (Rep (RuleM rep)))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp]
mergeexp)
            RuleM rep (Body rep -> MatchDec (BranchType rep) -> Exp rep)
-> RuleM rep (Body rep)
-> RuleM rep (MatchDec (BranchType rep) -> Exp rep)
forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Body rep -> RuleM rep (Body rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody Body rep
closedBody
            RuleM rep (MatchDec (BranchType rep) -> Exp rep)
-> RuleM rep (MatchDec (BranchType rep)) -> RuleM rep (Exp rep)
forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MatchDec (BranchType rep) -> RuleM rep (MatchDec (BranchType rep))
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([BranchType rep] -> MatchSort -> MatchDec (BranchType rep)
forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec [PrimType -> BranchType rep
forall rt. IsBodyType rt => PrimType -> rt
primBodyType PrimType
t] MatchSort
MatchNormal)
        )
  where
    ([FParam rep]
mergepat, [SubExp]
mergeexp) = [(FParam rep, SubExp)] -> ([FParam rep], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam rep, SubExp)]
merge
    mergeidents :: [Ident]
mergeidents = (FParam rep -> Ident) -> [FParam rep] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent [FParam rep]
mergepat
    mergenames :: [VName]
mergenames = (FParam rep -> VName) -> [FParam rep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> VName
forall dec. Param dec -> VName
paramName [FParam rep]
mergepat
    knownBnds :: Map VName SubExp
knownBnds = [(VName, SubExp)] -> Map VName SubExp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, SubExp)] -> Map VName SubExp)
-> [(VName, SubExp)] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ [VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
mergenames [SubExp]
mergeexp

checkResults ::
  (BuilderOps rep) =>
  [VName] ->
  SubExp ->
  Names ->
  IntType ->
  M.Map VName SubExp ->
  -- | Lambda-bound
  [VName] ->
  Body rep ->
  [SubExp] ->
  RuleM rep (Body rep)
checkResults :: forall rep.
BuilderOps rep =>
[VName]
-> SubExp
-> Names
-> IntType
-> Map VName SubExp
-> [VName]
-> Body rep
-> [SubExp]
-> RuleM rep (Body rep)
checkResults [VName]
pat SubExp
size Names
untouchable IntType
it Map VName SubExp
knownBnds [VName]
params Body rep
body [SubExp]
accs = do
  ((), Stms rep
stms) <-
    RuleM rep () -> RuleM rep ((), Stms (Rep (RuleM rep)))
forall a. RuleM rep a -> RuleM rep (a, Stms (Rep (RuleM rep)))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms (RuleM rep () -> RuleM rep ((), Stms (Rep (RuleM rep))))
-> RuleM rep () -> RuleM rep ((), Stms (Rep (RuleM rep)))
forall a b. (a -> b) -> a -> b
$
      ((VName, SubExpRes) -> (VName, SubExp) -> RuleM rep ())
-> [(VName, SubExpRes)] -> [(VName, SubExp)] -> RuleM rep ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (VName, SubExpRes) -> (VName, SubExp) -> RuleM rep ()
forall {rep}.
BuilderOps rep =>
(VName, SubExpRes) -> (VName, SubExp) -> RuleM rep ()
checkResult ([VName] -> [SubExpRes] -> [(VName, SubExpRes)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pat [SubExpRes]
res) ([VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
accparams [SubExp]
accs)
  Stms (Rep (RuleM rep))
-> [SubExpRes] -> RuleM rep (Body (Rep (RuleM rep)))
forall (m :: * -> *).
MonadBuilder m =>
Stms (Rep m) -> [SubExpRes] -> m (Body (Rep m))
mkBodyM Stms rep
Stms (Rep (RuleM rep))
stms ([SubExpRes] -> RuleM rep (Body (Rep (RuleM rep))))
-> [SubExpRes] -> RuleM rep (Body (Rep (RuleM rep)))
forall a b. (a -> b) -> a -> b
$ [VName] -> [SubExpRes]
varsRes [VName]
pat
  where
    stmMap :: Map VName (Exp rep)
stmMap = Body rep -> Map VName (Exp rep)
forall rep. Body rep -> Map VName (Exp rep)
makeBindMap Body rep
body
    ([VName]
accparams, [VName]
_) = Int -> [VName] -> ([VName], [VName])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
accs) [VName]
params
    res :: [SubExpRes]
res = Body rep -> [SubExpRes]
forall rep. Body rep -> [SubExpRes]
bodyResult Body rep
body

    nonFree :: Names
nonFree = Body rep -> Names
forall rep. Body rep -> Names
boundInBody Body rep
body Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList [VName]
params Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
untouchable

    checkResult :: (VName, SubExpRes) -> (VName, SubExp) -> RuleM rep ()
checkResult (VName
p, SubExpRes Certs
_ (Var VName
v)) (VName
accparam, SubExp
acc)
      | Just (BasicOp (BinOp BinOp
bop SubExp
x SubExp
y)) <- VName -> Map VName (Exp rep) -> Maybe (Exp rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName (Exp rep)
stmMap,
        SubExp
x SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
/= SubExp
y = do
          -- One of x,y must be *this* accumulator, and the other must
          -- be something that is free in the body.
          let isThisAccum :: SubExp -> Bool
isThisAccum = (SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
Var VName
accparam)
          (SubExp
this, SubExp
el) <- Maybe (SubExp, SubExp) -> RuleM rep (SubExp, SubExp)
forall a rep. Maybe a -> RuleM rep a
liftMaybe (Maybe (SubExp, SubExp) -> RuleM rep (SubExp, SubExp))
-> Maybe (SubExp, SubExp) -> RuleM rep (SubExp, SubExp)
forall a b. (a -> b) -> a -> b
$
            case ( (SubExp -> Maybe SubExp
asFreeSubExp SubExp
x, SubExp -> Bool
isThisAccum SubExp
y),
                   (SubExp -> Maybe SubExp
asFreeSubExp SubExp
y, SubExp -> Bool
isThisAccum SubExp
x)
                 ) of
              ((Just SubExp
free, Bool
True), (Maybe SubExp, Bool)
_) -> (SubExp, SubExp) -> Maybe (SubExp, SubExp)
forall a. a -> Maybe a
Just (SubExp
acc, SubExp
free)
              ((Maybe SubExp, Bool)
_, (Just SubExp
free, Bool
True)) -> (SubExp, SubExp) -> Maybe (SubExp, SubExp)
forall a. a -> Maybe a
Just (SubExp
acc, SubExp
free)
              ((Maybe SubExp, Bool), (Maybe SubExp, Bool))
_ -> Maybe (SubExp, SubExp)
forall a. Maybe a
Nothing

          case BinOp
bop of
            BinOp
LogAnd ->
              [VName] -> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
p] (Exp (Rep (RuleM rep)) -> RuleM rep ())
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (RuleM rep))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (RuleM rep)))
-> BasicOp -> Exp (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
LogAnd SubExp
this SubExp
el
            Add IntType
t Overflow
w -> do
              SubExp
size' <- IntType -> SubExp -> RuleM rep SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
t SubExp
size
              [VName] -> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
p]
                (Exp rep -> RuleM rep ()) -> RuleM rep (Exp rep) -> RuleM rep ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> RuleM rep (Exp (Rep (RuleM rep)))
-> RuleM rep (Exp (Rep (RuleM rep)))
-> RuleM rep (Exp (Rep (RuleM rep)))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
                  (IntType -> Overflow -> BinOp
Add IntType
t Overflow
w)
                  (SubExp -> RuleM rep (Exp (Rep (RuleM rep)))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
this)
                  (Exp (Rep (RuleM rep)) -> RuleM rep (Exp (Rep (RuleM rep)))
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep (RuleM rep)) -> RuleM rep (Exp (Rep (RuleM rep))))
-> Exp (Rep (RuleM rep)) -> RuleM rep (Exp (Rep (RuleM rep)))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (RuleM rep))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (RuleM rep)))
-> BasicOp -> Exp (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Overflow -> BinOp
Mul IntType
t Overflow
w) SubExp
el SubExp
size')
            FAdd FloatType
t | Just RuleM rep SubExp
properly_typed_size <- FloatType -> Maybe (RuleM rep SubExp)
forall {m :: * -> *}.
MonadBuilder m =>
FloatType -> Maybe (m SubExp)
properFloatSize FloatType
t -> do
              SubExp
size' <- RuleM rep SubExp
properly_typed_size
              [VName] -> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
p]
                (Exp rep -> RuleM rep ()) -> RuleM rep (Exp rep) -> RuleM rep ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> RuleM rep (Exp (Rep (RuleM rep)))
-> RuleM rep (Exp (Rep (RuleM rep)))
-> RuleM rep (Exp (Rep (RuleM rep)))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
                  (FloatType -> BinOp
FAdd FloatType
t)
                  (SubExp -> RuleM rep (Exp (Rep (RuleM rep)))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
this)
                  (Exp (Rep (RuleM rep)) -> RuleM rep (Exp (Rep (RuleM rep)))
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep (RuleM rep)) -> RuleM rep (Exp (Rep (RuleM rep))))
-> Exp (Rep (RuleM rep)) -> RuleM rep (Exp (Rep (RuleM rep)))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (RuleM rep))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (RuleM rep)))
-> BasicOp -> Exp (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (FloatType -> BinOp
FMul FloatType
t) SubExp
el SubExp
size')
            BinOp
_ -> RuleM rep ()
forall rep a. RuleM rep a
cannotSimplify -- Um... sorry.
    checkResult (VName, SubExpRes)
_ (VName, SubExp)
_ = RuleM rep ()
forall rep a. RuleM rep a
cannotSimplify

    asFreeSubExp :: SubExp -> Maybe SubExp
    asFreeSubExp :: SubExp -> Maybe SubExp
asFreeSubExp (Var VName
v)
      | VName
v VName -> Names -> Bool
`nameIn` Names
nonFree = VName -> Map VName SubExp -> Maybe SubExp
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName SubExp
knownBnds
    asFreeSubExp SubExp
se = SubExp -> Maybe SubExp
forall a. a -> Maybe a
Just SubExp
se

    properFloatSize :: FloatType -> Maybe (m SubExp)
properFloatSize FloatType
t =
      m SubExp -> Maybe (m SubExp)
forall a. a -> Maybe a
Just (m SubExp -> Maybe (m SubExp)) -> m SubExp -> Maybe (m SubExp)
forall a b. (a -> b) -> a -> b
$
        String -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"converted_size" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
            ConvOp -> SubExp -> BasicOp
ConvOp (IntType -> FloatType -> ConvOp
SIToFP IntType
it FloatType
t) SubExp
size

determineKnownBindings ::
  VarLookup rep ->
  Lambda rep ->
  [SubExp] ->
  [VName] ->
  M.Map VName SubExp
determineKnownBindings :: forall rep.
VarLookup rep
-> Lambda rep -> [SubExp] -> [VName] -> Map VName SubExp
determineKnownBindings VarLookup rep
look Lambda rep
lam [SubExp]
accs [VName]
arrs =
  Map VName SubExp
accBnds Map VName SubExp -> Map VName SubExp -> Map VName SubExp
forall a. Semigroup a => a -> a -> a
<> Map VName SubExp
arrBnds
  where
    ([Param (LParamInfo rep)]
accparams, [Param (LParamInfo rep)]
arrparams) =
      Int
-> [Param (LParamInfo rep)]
-> ([Param (LParamInfo rep)], [Param (LParamInfo rep)])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
accs) ([Param (LParamInfo rep)]
 -> ([Param (LParamInfo rep)], [Param (LParamInfo rep)]))
-> [Param (LParamInfo rep)]
-> ([Param (LParamInfo rep)], [Param (LParamInfo rep)])
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Param (LParamInfo rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
    accBnds :: Map VName SubExp
accBnds =
      [(VName, SubExp)] -> Map VName SubExp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, SubExp)] -> Map VName SubExp)
-> [(VName, SubExp)] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$
        [VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
accparams) [SubExp]
accs
    arrBnds :: Map VName SubExp
arrBnds =
      [(VName, SubExp)] -> Map VName SubExp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, SubExp)] -> Map VName SubExp)
-> [(VName, SubExp)] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$
        ((VName, VName) -> Maybe (VName, SubExp))
-> [(VName, VName)] -> [(VName, SubExp)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName, VName) -> Maybe (VName, SubExp)
forall {a}. (a, VName) -> Maybe (a, SubExp)
isReplicate ([(VName, VName)] -> [(VName, SubExp)])
-> [(VName, VName)] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$
          [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
arrparams) [VName]
arrs

    isReplicate :: (a, VName) -> Maybe (a, SubExp)
isReplicate (a
p, VName
v)
      | Just (BasicOp (Replicate (Shape (SubExp
_ : [SubExp]
_)) SubExp
ve), Certs
cs) <- VarLookup rep
look VName
v,
        Certs
cs Certs -> Certs -> Bool
forall a. Eq a => a -> a -> Bool
== Certs
forall a. Monoid a => a
mempty =
          (a, SubExp) -> Maybe (a, SubExp)
forall a. a -> Maybe a
Just (a
p, SubExp
ve)
    isReplicate (a, VName)
_ = Maybe (a, SubExp)
forall a. Maybe a
Nothing

makeBindMap :: Body rep -> M.Map VName (Exp rep)
makeBindMap :: forall rep. Body rep -> Map VName (Exp rep)
makeBindMap = [(VName, Exp rep)] -> Map VName (Exp rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Exp rep)] -> Map VName (Exp rep))
-> (Body rep -> [(VName, Exp rep)])
-> Body rep
-> Map VName (Exp rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm rep -> Maybe (VName, Exp rep))
-> [Stm rep] -> [(VName, Exp rep)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Stm rep -> Maybe (VName, Exp rep)
forall {rep}. Stm rep -> Maybe (VName, Exp rep)
isSingletonStm ([Stm rep] -> [(VName, Exp rep)])
-> (Body rep -> [Stm rep]) -> Body rep -> [(VName, Exp rep)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms rep -> [Stm rep])
-> (Body rep -> Stms rep) -> Body rep -> [Stm rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms
  where
    isSingletonStm :: Stm rep -> Maybe (VName, Exp rep)
isSingletonStm (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ Exp rep
e) = case Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat of
      [VName
v] -> (VName, Exp rep) -> Maybe (VName, Exp rep)
forall a. a -> Maybe a
Just (VName
v, Exp rep
e)
      [VName]
_ -> Maybe (VName, Exp rep)
forall a. Maybe a
Nothing