{-# LANGUAGE Strict #-}

-- | Perform horizontal and vertical fusion of SOACs.  See the paper
-- /A T2 Graph-Reduction Approach To Fusion/ for the basic idea (some
-- extensions discussed in /Design and GPGPU Performance of Futhark’s
-- Redomap Construct/).
module Futhark.Optimise.Fusion (fuseSOACs) where

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Data.Graph.Inductive.Graph qualified as G
import Data.Graph.Inductive.Query.DFS qualified as Q
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.Alias qualified as Alias
import Futhark.Analysis.HORep.SOAC qualified as H
import Futhark.Construct
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS hiding (SOAC (..))
import Futhark.IR.SOACS qualified as Futhark
import Futhark.IR.SOACS.Simplify (simplifyLambda)
import Futhark.Optimise.Fusion.GraphRep
import Futhark.Optimise.Fusion.RulesWithAccs qualified as SF
import Futhark.Optimise.Fusion.TryFusion qualified as TF
import Futhark.Pass
import Futhark.Transform.Rename
import Futhark.Transform.Substitute

data FusionEnv = FusionEnv
  { FusionEnv -> VNameSource
vNameSource :: VNameSource,
    FusionEnv -> Int
fusionCount :: Int,
    FusionEnv -> Bool
fuseScans :: Bool,
    FusionEnv -> Maybe Int
gas :: Maybe Int
  }

freshFusionEnv :: Maybe Int -> FusionEnv
freshFusionEnv :: Maybe Int -> FusionEnv
freshFusionEnv Maybe Int
gas =
  FusionEnv
    { vNameSource :: VNameSource
vNameSource = VNameSource
blankNameSource,
      fusionCount :: Int
fusionCount = Int
0,
      fuseScans :: Bool
fuseScans = Bool
True,
      gas :: Maybe Int
gas = Maybe Int
gas
    }

newtype FusionM a = FusionM (ReaderT (Scope SOACS) (State FusionEnv) a)
  deriving
    ( Applicative FusionM
Applicative FusionM =>
(forall a b. FusionM a -> (a -> FusionM b) -> FusionM b)
-> (forall a b. FusionM a -> FusionM b -> FusionM b)
-> (forall a. a -> FusionM a)
-> Monad FusionM
forall a. a -> FusionM a
forall a b. FusionM a -> FusionM b -> FusionM b
forall a b. FusionM a -> (a -> FusionM b) -> FusionM b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. FusionM a -> (a -> FusionM b) -> FusionM b
>>= :: forall a b. FusionM a -> (a -> FusionM b) -> FusionM b
$c>> :: forall a b. FusionM a -> FusionM b -> FusionM b
>> :: forall a b. FusionM a -> FusionM b -> FusionM b
$creturn :: forall a. a -> FusionM a
return :: forall a. a -> FusionM a
Monad,
      Functor FusionM
Functor FusionM =>
(forall a. a -> FusionM a)
-> (forall a b. FusionM (a -> b) -> FusionM a -> FusionM b)
-> (forall a b c.
    (a -> b -> c) -> FusionM a -> FusionM b -> FusionM c)
-> (forall a b. FusionM a -> FusionM b -> FusionM b)
-> (forall a b. FusionM a -> FusionM b -> FusionM a)
-> Applicative FusionM
forall a. a -> FusionM a
forall a b. FusionM a -> FusionM b -> FusionM a
forall a b. FusionM a -> FusionM b -> FusionM b
forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
forall a b c. (a -> b -> c) -> FusionM a -> FusionM b -> FusionM c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> FusionM a
pure :: forall a. a -> FusionM a
$c<*> :: forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
<*> :: forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
$cliftA2 :: forall a b c. (a -> b -> c) -> FusionM a -> FusionM b -> FusionM c
liftA2 :: forall a b c. (a -> b -> c) -> FusionM a -> FusionM b -> FusionM c
$c*> :: forall a b. FusionM a -> FusionM b -> FusionM b
*> :: forall a b. FusionM a -> FusionM b -> FusionM b
$c<* :: forall a b. FusionM a -> FusionM b -> FusionM a
<* :: forall a b. FusionM a -> FusionM b -> FusionM a
Applicative,
      (forall a b. (a -> b) -> FusionM a -> FusionM b)
-> (forall a b. a -> FusionM b -> FusionM a) -> Functor FusionM
forall a b. a -> FusionM b -> FusionM a
forall a b. (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> FusionM a -> FusionM b
fmap :: forall a b. (a -> b) -> FusionM a -> FusionM b
$c<$ :: forall a b. a -> FusionM b -> FusionM a
<$ :: forall a b. a -> FusionM b -> FusionM a
Functor,
      MonadState FusionEnv,
      HasScope SOACS,
      LocalScope SOACS
    )

instance MonadFreshNames FusionM where
  getNameSource :: FusionM VNameSource
getNameSource = (FusionEnv -> VNameSource) -> FusionM VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets FusionEnv -> VNameSource
vNameSource
  putNameSource :: VNameSource -> FusionM ()
putNameSource VNameSource
source =
    (FusionEnv -> FusionEnv) -> FusionM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\FusionEnv
env -> FusionEnv
env {vNameSource = source})

useGas :: a -> FusionM a -> FusionM a
useGas :: forall a. a -> FusionM a -> FusionM a
useGas a
g FusionM a
m = do
  is_out_of_gas <- (FusionEnv -> Bool) -> FusionM Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Bool -> (Int -> Bool) -> Maybe Int -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) (Maybe Int -> Bool)
-> (FusionEnv -> Maybe Int) -> FusionEnv -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FusionEnv -> Maybe Int
gas)
  if is_out_of_gas
    then pure g
    else
      modify (\FusionEnv
s -> FusionEnv
s {gas = (\Int
gas' -> Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
gas' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) <$> gas s})
        >> m

runFusionM :: (MonadFreshNames m) => Scope SOACS -> FusionEnv -> FusionM a -> m a
runFusionM :: forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> FusionEnv -> FusionM a -> m a
runFusionM Scope SOACS
scope FusionEnv
fenv (FusionM ReaderT (Scope SOACS) (State FusionEnv) a
a) = (VNameSource -> (a, VNameSource)) -> m a
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (a, VNameSource)) -> m a)
-> (VNameSource -> (a, VNameSource)) -> m a
forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
  let x :: State FusionEnv a
x = ReaderT (Scope SOACS) (State FusionEnv) a
-> Scope SOACS -> State FusionEnv a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Scope SOACS) (State FusionEnv) a
a Scope SOACS
scope
      (a
y, FusionEnv
z) = State FusionEnv a -> FusionEnv -> (a, FusionEnv)
forall s a. State s a -> s -> (a, s)
runState State FusionEnv a
x (FusionEnv
fenv {vNameSource = src})
   in (a
y, FusionEnv -> VNameSource
vNameSource FusionEnv
z)

doFuseScans :: FusionM a -> FusionM a
doFuseScans :: forall a. FusionM a -> FusionM a
doFuseScans FusionM a
m = do
  fs <- (FusionEnv -> Bool) -> FusionM Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets FusionEnv -> Bool
fuseScans
  modify (\FusionEnv
s -> FusionEnv
s {fuseScans = True})
  r <- m
  modify (\FusionEnv
s -> FusionEnv
s {fuseScans = fs})
  pure r

dontFuseScans :: FusionM a -> FusionM a
dontFuseScans :: forall a. FusionM a -> FusionM a
dontFuseScans FusionM a
m = do
  fs <- (FusionEnv -> Bool) -> FusionM Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets FusionEnv -> Bool
fuseScans
  modify (\FusionEnv
s -> FusionEnv
s {fuseScans = False})
  r <- m
  modify (\FusionEnv
s -> FusionEnv
s {fuseScans = fs})
  pure r

isNotVarInput :: [H.Input] -> [H.Input]
isNotVarInput :: [Input] -> [Input]
isNotVarInput = (Input -> Bool) -> [Input] -> [Input]
forall a. (a -> Bool) -> [a] -> [a]
filter (Maybe VName -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe VName -> Bool) -> (Input -> Maybe VName) -> Input -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> Maybe VName
H.isVarInput)

finalizeNode :: (HasScope SOACS m, MonadFreshNames m) => NodeT -> m (Stms SOACS)
finalizeNode :: forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
NodeT -> m (Stms SOACS)
finalizeNode NodeT
nt = case NodeT
nt of
  StmNode Stm SOACS
stm -> Stms SOACS -> m (Stms SOACS)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS -> m (Stms SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm
  SoacNode ArrayTransforms
ots Pat Type
outputs SOAC SOACS
soac StmAux (ExpDec SOACS)
aux -> Builder SOACS () -> m (Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (Stms rep)
runBuilder_ (Builder SOACS () -> m (Stms SOACS))
-> Builder SOACS () -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ do
    untransformed_outputs <- (VName -> BuilderT SOACS (State VNameSource) VName)
-> [VName] -> BuilderT SOACS (State VNameSource) [VName]
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 -> BuilderT SOACS (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName ([VName] -> BuilderT SOACS (State VNameSource) [VName])
-> [VName] -> BuilderT SOACS (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
outputs
    auxing aux $ letBindNames untransformed_outputs . Op =<< H.toSOAC soac
    forM_ (zip (patNames outputs) untransformed_outputs) $ \(VName
output, VName
v) ->
      [VName]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> Builder SOACS ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
output] (Exp SOACS -> Builder SOACS ())
-> (VName -> Exp SOACS) -> VName -> Builder SOACS ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> (VName -> BasicOp) -> VName -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var (VName -> Builder SOACS ())
-> BuilderT SOACS (State VNameSource) VName -> Builder SOACS ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ArrayTransforms
-> VName -> BuilderT SOACS (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
ArrayTransforms -> VName -> m VName
H.applyTransforms ArrayTransforms
ots VName
v
  ResNode VName
_ -> Stms SOACS -> m (Stms SOACS)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stms SOACS
forall a. Monoid a => a
mempty
  TransNode VName
output ArrayTransform
tr VName
ia -> do
    (aux, e) <- ArrayTransform -> VName -> m (StmAux (), Exp SOACS)
forall (m :: * -> *) rep.
(Monad m, HasScope rep m) =>
ArrayTransform -> VName -> m (StmAux (), Exp rep)
H.transformToExp ArrayTransform
tr VName
ia
    runBuilder_ $ auxing aux $ letBindNames [output] e
  FreeNode VName
_ -> Stms SOACS -> m (Stms SOACS)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stms SOACS
forall a. Monoid a => a
mempty
  DoNode Stm SOACS
stm [(NodeT, [EdgeT])]
lst -> do
    lst' <- ((NodeT, [EdgeT]) -> m (Stms SOACS))
-> [(NodeT, [EdgeT])] -> m [Stms SOACS]
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 (NodeT -> m (Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
NodeT -> m (Stms SOACS)
finalizeNode (NodeT -> m (Stms SOACS))
-> ((NodeT, [EdgeT]) -> NodeT)
-> (NodeT, [EdgeT])
-> m (Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NodeT, [EdgeT]) -> NodeT
forall a b. (a, b) -> a
fst) [(NodeT, [EdgeT])]
lst
    pure $ mconcat lst' <> oneStm stm
  MatchNode Stm SOACS
stm [(NodeT, [EdgeT])]
lst -> do
    lst' <- ((NodeT, [EdgeT]) -> m (Stms SOACS))
-> [(NodeT, [EdgeT])] -> m [Stms SOACS]
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 (NodeT -> m (Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
NodeT -> m (Stms SOACS)
finalizeNode (NodeT -> m (Stms SOACS))
-> ((NodeT, [EdgeT]) -> NodeT)
-> (NodeT, [EdgeT])
-> m (Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NodeT, [EdgeT]) -> NodeT
forall a b. (a, b) -> a
fst) [(NodeT, [EdgeT])]
lst
    pure $ mconcat lst' <> oneStm stm

linearizeGraph :: (HasScope SOACS m, MonadFreshNames m) => DepGraph -> m (Stms SOACS)
linearizeGraph :: forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
DepGraph -> m (Stms SOACS)
linearizeGraph DepGraph
dg =
  ([Stms SOACS] -> Stms SOACS) -> m [Stms SOACS] -> m (Stms SOACS)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Stms SOACS] -> Stms SOACS
forall a. Monoid a => [a] -> a
mconcat (m [Stms SOACS] -> m (Stms SOACS))
-> m [Stms SOACS] -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ (NodeT -> m (Stms SOACS)) -> [NodeT] -> m [Stms SOACS]
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 NodeT -> m (Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
NodeT -> m (Stms SOACS)
finalizeNode ([NodeT] -> m [Stms SOACS]) -> [NodeT] -> m [Stms SOACS]
forall a b. (a -> b) -> a -> b
$ [NodeT] -> [NodeT]
forall a. [a] -> [a]
reverse ([NodeT] -> [NodeT]) -> [NodeT] -> [NodeT]
forall a b. (a -> b) -> a -> b
$ Gr NodeT EdgeT -> [NodeT]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [a]
Q.topsort' (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg)

fusedSomething :: NodeT -> FusionM (Maybe NodeT)
fusedSomething :: NodeT -> FusionM (Maybe NodeT)
fusedSomething NodeT
x = do
  (FusionEnv -> FusionEnv) -> FusionM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FusionEnv -> FusionEnv) -> FusionM ())
-> (FusionEnv -> FusionEnv) -> FusionM ()
forall a b. (a -> b) -> a -> b
$ \FusionEnv
s -> FusionEnv
s {fusionCount = 1 + fusionCount s}
  Maybe NodeT -> FusionM (Maybe NodeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe NodeT -> FusionM (Maybe NodeT))
-> Maybe NodeT -> FusionM (Maybe NodeT)
forall a b. (a -> b) -> a -> b
$ NodeT -> Maybe NodeT
forall a. a -> Maybe a
Just NodeT
x

vTryFuseNodesInGraph :: G.Node -> G.Node -> DepGraphAug FusionM
-- find the neighbors -> verify that fusion causes no cycles -> fuse
vTryFuseNodesInGraph :: Int -> Int -> DepGraphAug FusionM
vTryFuseNodesInGraph Int
node_1 Int
node_2 dg :: DepGraph
dg@DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
dgGraph = Gr NodeT EdgeT
g}
  | Bool -> Bool
not (Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
node_1 Gr NodeT EdgeT
g Bool -> Bool -> Bool
&& Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
node_2 Gr NodeT EdgeT
g) = DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  | DepGraph -> Int -> Int -> Bool
vFusionFeasability DepGraph
dg Int
node_1 Int
node_2 =
      DepGraph -> FusionM DepGraph -> FusionM DepGraph
forall a. a -> FusionM a -> FusionM a
useGas DepGraph
dg (FusionM DepGraph -> FusionM DepGraph)
-> FusionM DepGraph -> FusionM DepGraph
forall a b. (a -> b) -> a -> b
$ do
        let (Context NodeT EdgeT
ctx1, Context NodeT EdgeT
ctx2) = (Gr NodeT EdgeT -> Int -> Context NodeT EdgeT
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Context a b
G.context Gr NodeT EdgeT
g Int
node_1, Gr NodeT EdgeT -> Int -> Context NodeT EdgeT
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Context a b
G.context Gr NodeT EdgeT
g Int
node_2)
        fres <- [EdgeT]
-> [VName]
-> Context NodeT EdgeT
-> Context NodeT EdgeT
-> FusionM (Maybe (Context NodeT EdgeT))
vFuseContexts [EdgeT]
edgs [VName]
infusable_nodes Context NodeT EdgeT
ctx1 Context NodeT EdgeT
ctx2
        case fres of
          Just (Adj EdgeT
inputs, Int
_, NodeT
nodeT, Adj EdgeT
outputs) -> do
            nodeT' <-
              if [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
fusedC
                then NodeT -> FusionM NodeT
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodeT
nodeT
                else do
                  let (Adj EdgeT
_, Int
_, NodeT
_, Adj EdgeT
deps_1) = Context NodeT EdgeT
ctx1
                      (Adj EdgeT
_, Int
_, NodeT
_, Adj EdgeT
deps_2) = Context NodeT EdgeT
ctx2
                      -- make copies of everything that was not
                      -- previously consumed
                      old_cons :: [VName]
old_cons = ((EdgeT, Int) -> VName) -> Adj EdgeT -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (EdgeT -> VName
getName (EdgeT -> VName)
-> ((EdgeT, Int) -> EdgeT) -> (EdgeT, Int) -> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EdgeT, Int) -> EdgeT
forall a b. (a, b) -> a
fst) (Adj EdgeT -> [VName]) -> Adj EdgeT -> [VName]
forall a b. (a -> b) -> a -> b
$ ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
forall a. (a -> Bool) -> [a] -> [a]
filter (EdgeT -> Bool
isCons (EdgeT -> Bool) -> ((EdgeT, Int) -> EdgeT) -> (EdgeT, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EdgeT, Int) -> EdgeT
forall a b. (a, b) -> a
fst) (Adj EdgeT
deps_1 Adj EdgeT -> Adj EdgeT -> Adj EdgeT
forall a. Semigroup a => a -> a -> a
<> Adj EdgeT
deps_2)
                  [VName] -> NodeT -> FusionM NodeT
forall (m :: * -> *).
(LocalScope SOACS m, MonadFreshNames m) =>
[VName] -> NodeT -> m NodeT
makeCopiesOfFusedExcept [VName]
old_cons NodeT
nodeT
            contractEdge node_2 (inputs, node_1, nodeT', outputs) dg
          Maybe (Context NodeT EdgeT)
Nothing -> DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  | Bool
otherwise = DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  where
    edgs :: [EdgeT]
edgs = (DepEdge -> EdgeT) -> [DepEdge] -> [EdgeT]
forall a b. (a -> b) -> [a] -> [b]
map DepEdge -> EdgeT
forall b. LEdge b -> b
G.edgeLabel ([DepEdge] -> [EdgeT]) -> [DepEdge] -> [EdgeT]
forall a b. (a -> b) -> a -> b
$ DepGraph -> Int -> Int -> [DepEdge]
edgesBetween DepGraph
dg Int
node_1 Int
node_2
    fusedC :: [VName]
fusedC = (EdgeT -> VName) -> [EdgeT] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map EdgeT -> VName
getName ([EdgeT] -> [VName]) -> [EdgeT] -> [VName]
forall a b. (a -> b) -> a -> b
$ (EdgeT -> Bool) -> [EdgeT] -> [EdgeT]
forall a. (a -> Bool) -> [a] -> [a]
filter EdgeT -> Bool
isCons [EdgeT]
edgs
    infusable_nodes :: [VName]
infusable_nodes =
      (DepEdge -> VName) -> [DepEdge] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map
        DepEdge -> VName
depsFromEdge
        ((Int -> [DepEdge]) -> [Int] -> [DepEdge]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (DepGraph -> Int -> Int -> [DepEdge]
edgesBetween DepGraph
dg Int
node_1) ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
node_2) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Gr NodeT EdgeT -> Int -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> Int -> [Int]
G.pre Gr NodeT EdgeT
g Int
node_1))

hTryFuseNodesInGraph :: G.Node -> G.Node -> DepGraphAug FusionM
hTryFuseNodesInGraph :: Int -> Int -> DepGraphAug FusionM
hTryFuseNodesInGraph Int
node_1 Int
node_2 dg :: DepGraph
dg@DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
dgGraph = Gr NodeT EdgeT
g}
  | Bool -> Bool
not (Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
node_1 Gr NodeT EdgeT
g Bool -> Bool -> Bool
&& Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
node_2 Gr NodeT EdgeT
g) = DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  | DepGraph -> Int -> Int -> Bool
hFusionFeasability DepGraph
dg Int
node_1 Int
node_2 =
      DepGraph -> FusionM DepGraph -> FusionM DepGraph
forall a. a -> FusionM a -> FusionM a
useGas DepGraph
dg (FusionM DepGraph -> FusionM DepGraph)
-> FusionM DepGraph -> FusionM DepGraph
forall a b. (a -> b) -> a -> b
$ do
        fres <- Context NodeT EdgeT
-> Context NodeT EdgeT -> FusionM (Maybe (Context NodeT EdgeT))
hFuseContexts (Gr NodeT EdgeT -> Int -> Context NodeT EdgeT
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Context a b
G.context Gr NodeT EdgeT
g Int
node_1) (Gr NodeT EdgeT -> Int -> Context NodeT EdgeT
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Context a b
G.context Gr NodeT EdgeT
g Int
node_2)
        case fres of
          Just Context NodeT EdgeT
ctx -> Int -> Context NodeT EdgeT -> DepGraphAug FusionM
forall (m :: * -> *).
Monad m =>
Int -> Context NodeT EdgeT -> DepGraphAug m
contractEdge Int
node_2 Context NodeT EdgeT
ctx DepGraph
dg
          Maybe (Context NodeT EdgeT)
Nothing -> DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  | Bool
otherwise = DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg

hFuseContexts :: DepContext -> DepContext -> FusionM (Maybe DepContext)
hFuseContexts :: Context NodeT EdgeT
-> Context NodeT EdgeT -> FusionM (Maybe (Context NodeT EdgeT))
hFuseContexts Context NodeT EdgeT
c1 Context NodeT EdgeT
c2 = do
  let (Adj EdgeT
_, Int
_, NodeT
nodeT1, Adj EdgeT
_) = Context NodeT EdgeT
c1
      (Adj EdgeT
_, Int
_, NodeT
nodeT2, Adj EdgeT
_) = Context NodeT EdgeT
c2
  fres <- NodeT -> NodeT -> FusionM (Maybe NodeT)
hFuseNodeT NodeT
nodeT1 NodeT
nodeT2
  case fres of
    Just NodeT
nodeT -> Maybe (Context NodeT EdgeT)
-> FusionM (Maybe (Context NodeT EdgeT))
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Context NodeT EdgeT)
 -> FusionM (Maybe (Context NodeT EdgeT)))
-> Maybe (Context NodeT EdgeT)
-> FusionM (Maybe (Context NodeT EdgeT))
forall a b. (a -> b) -> a -> b
$ Context NodeT EdgeT -> Maybe (Context NodeT EdgeT)
forall a. a -> Maybe a
Just (NodeT
-> Context NodeT EdgeT
-> Context NodeT EdgeT
-> Context NodeT EdgeT
forall b a. Ord b => a -> Context a b -> Context a b -> Context a b
mergedContext NodeT
nodeT Context NodeT EdgeT
c1 Context NodeT EdgeT
c2)
    Maybe NodeT
Nothing -> Maybe (Context NodeT EdgeT)
-> FusionM (Maybe (Context NodeT EdgeT))
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Context NodeT EdgeT)
forall a. Maybe a
Nothing

vFuseContexts :: [EdgeT] -> [VName] -> DepContext -> DepContext -> FusionM (Maybe DepContext)
vFuseContexts :: [EdgeT]
-> [VName]
-> Context NodeT EdgeT
-> Context NodeT EdgeT
-> FusionM (Maybe (Context NodeT EdgeT))
vFuseContexts [EdgeT]
edgs [VName]
infusable Context NodeT EdgeT
c1 Context NodeT EdgeT
c2 = do
  let (Adj EdgeT
i1, Int
n1, NodeT
nodeT1, Adj EdgeT
o1) = Context NodeT EdgeT
c1
      (Adj EdgeT
_i2, Int
n2, NodeT
nodeT2, Adj EdgeT
o2) = Context NodeT EdgeT
c2
  fres <-
    [EdgeT]
-> [VName]
-> (NodeT, [EdgeT], [EdgeT])
-> (NodeT, [EdgeT])
-> FusionM (Maybe NodeT)
vFuseNodeT
      [EdgeT]
edgs
      [VName]
infusable
      (NodeT
nodeT1, ((EdgeT, Int) -> EdgeT) -> Adj EdgeT -> [EdgeT]
forall a b. (a -> b) -> [a] -> [b]
map (EdgeT, Int) -> EdgeT
forall a b. (a, b) -> a
fst (Adj EdgeT -> [EdgeT]) -> Adj EdgeT -> [EdgeT]
forall a b. (a -> b) -> a -> b
$ ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(/=) Int
n2 (Int -> Bool) -> ((EdgeT, Int) -> Int) -> (EdgeT, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EdgeT, Int) -> Int
forall a b. (a, b) -> b
snd) Adj EdgeT
i1, ((EdgeT, Int) -> EdgeT) -> Adj EdgeT -> [EdgeT]
forall a b. (a -> b) -> [a] -> [b]
map (EdgeT, Int) -> EdgeT
forall a b. (a, b) -> a
fst Adj EdgeT
o1)
      (NodeT
nodeT2, ((EdgeT, Int) -> EdgeT) -> Adj EdgeT -> [EdgeT]
forall a b. (a -> b) -> [a] -> [b]
map (EdgeT, Int) -> EdgeT
forall a b. (a, b) -> a
fst (Adj EdgeT -> [EdgeT]) -> Adj EdgeT -> [EdgeT]
forall a b. (a -> b) -> a -> b
$ ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(/=) Int
n1 (Int -> Bool) -> ((EdgeT, Int) -> Int) -> (EdgeT, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EdgeT, Int) -> Int
forall a b. (a, b) -> b
snd) Adj EdgeT
o2)
  case fres of
    Just NodeT
nodeT -> Maybe (Context NodeT EdgeT)
-> FusionM (Maybe (Context NodeT EdgeT))
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Context NodeT EdgeT)
 -> FusionM (Maybe (Context NodeT EdgeT)))
-> Maybe (Context NodeT EdgeT)
-> FusionM (Maybe (Context NodeT EdgeT))
forall a b. (a -> b) -> a -> b
$ Context NodeT EdgeT -> Maybe (Context NodeT EdgeT)
forall a. a -> Maybe a
Just (NodeT
-> Context NodeT EdgeT
-> Context NodeT EdgeT
-> Context NodeT EdgeT
forall b a. Ord b => a -> Context a b -> Context a b -> Context a b
mergedContext NodeT
nodeT Context NodeT EdgeT
c1 Context NodeT EdgeT
c2)
    Maybe NodeT
Nothing -> Maybe (Context NodeT EdgeT)
-> FusionM (Maybe (Context NodeT EdgeT))
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Context NodeT EdgeT)
forall a. Maybe a
Nothing

makeCopiesOfFusedExcept ::
  (LocalScope SOACS m, MonadFreshNames m) =>
  [VName] ->
  NodeT ->
  m NodeT
makeCopiesOfFusedExcept :: forall (m :: * -> *).
(LocalScope SOACS m, MonadFreshNames m) =>
[VName] -> NodeT -> m NodeT
makeCopiesOfFusedExcept [VName]
noCopy (SoacNode ArrayTransforms
ots Pat Type
pats SOAC SOACS
soac StmAux (ExpDec SOACS)
aux) = do
  case SOAC SOACS
soac of
    H.Screma SubExp
w [Input]
arrs (ScremaForm Lambda SOACS
lam [Scan SOACS]
scans [Reduce SOACS]
reduces Lambda SOACS
postlam) -> do
      Scope SOACS -> m NodeT -> m NodeT
forall a. Scope SOACS -> m a -> m a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Lambda SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Lambda SOACS
lam Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Lambda SOACS
postlam) (m NodeT -> m NodeT) -> m NodeT -> m NodeT
forall a b. (a -> b) -> a -> b
$ do
        fused_in_main <-
          (VName -> m Bool) -> [VName] -> m [VName]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM ((Type -> Bool) -> m Type -> m Bool
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> Bool
not (Bool -> Bool) -> (Type -> Bool) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc) (m Type -> m Bool) -> (VName -> m Type) -> VName -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType) ([VName] -> m [VName])
-> (Lambda (Aliases SOACS) -> [VName])
-> Lambda (Aliases SOACS)
-> m [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList (Names -> [VName])
-> (Lambda (Aliases SOACS) -> Names)
-> Lambda (Aliases SOACS)
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda (Aliases SOACS) -> Names
forall rep. Aliased rep => Lambda rep -> Names
consumedByLambda (Lambda (Aliases SOACS) -> m [VName])
-> Lambda (Aliases SOACS) -> m [VName]
forall a b. (a -> b) -> a -> b
$
            AliasTable -> Lambda SOACS -> Lambda (Aliases SOACS)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
forall a. Monoid a => a
mempty Lambda SOACS
lam
        fused_in_post <-
          filterM (fmap (not . isAcc) . lookupType) . namesToList . consumedByLambda $
            Alias.analyseLambda mempty postlam
        lam' <- makeCopiesInLambda (fused_in_main L.\\ noCopy) lam
        postlam' <- makeCopiesInLambda (fused_in_post L.\\ noCopy) postlam
        let form' = Lambda SOACS
-> [Scan SOACS]
-> [Reduce SOACS]
-> Lambda SOACS
-> ScremaForm SOACS
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm Lambda SOACS
lam' [Scan SOACS]
scans [Reduce SOACS]
reduces Lambda SOACS
postlam'
        pure $ SoacNode ots pats (H.Screma w arrs form') aux
    SOAC SOACS
_any -> do
      let lam :: Lambda SOACS
lam = SOAC SOACS -> Lambda SOACS
forall rep. SOAC rep -> Lambda rep
H.lambda SOAC SOACS
soac
      Scope SOACS -> m NodeT -> m NodeT
forall a. Scope SOACS -> m a -> m a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Lambda SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Lambda SOACS
lam) (m NodeT -> m NodeT) -> m NodeT -> m NodeT
forall a b. (a -> b) -> a -> b
$ do
        fused_inner <-
          (VName -> m Bool) -> [VName] -> m [VName]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM ((Type -> Bool) -> m Type -> m Bool
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> Bool
not (Bool -> Bool) -> (Type -> Bool) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc) (m Type -> m Bool) -> (VName -> m Type) -> VName -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType) ([VName] -> m [VName])
-> (Lambda (Aliases SOACS) -> [VName])
-> Lambda (Aliases SOACS)
-> m [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList (Names -> [VName])
-> (Lambda (Aliases SOACS) -> Names)
-> Lambda (Aliases SOACS)
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda (Aliases SOACS) -> Names
forall rep. Aliased rep => Lambda rep -> Names
consumedByLambda (Lambda (Aliases SOACS) -> m [VName])
-> Lambda (Aliases SOACS) -> m [VName]
forall a b. (a -> b) -> a -> b
$
            AliasTable -> Lambda SOACS -> Lambda (Aliases SOACS)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
forall a. Monoid a => a
mempty Lambda SOACS
lam
        lam' <- makeCopiesInLambda (fused_inner L.\\ noCopy) lam
        pure $ SoacNode ots pats (H.setLambda lam' soac) aux
makeCopiesOfFusedExcept [VName]
_ NodeT
nodeT = NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodeT
nodeT

makeCopiesInLambda ::
  (LocalScope SOACS m, MonadFreshNames m) =>
  [VName] ->
  Lambda SOACS ->
  m (Lambda SOACS)
makeCopiesInLambda :: forall (m :: * -> *).
(LocalScope SOACS m, MonadFreshNames m) =>
[VName] -> Lambda SOACS -> m (Lambda SOACS)
makeCopiesInLambda [VName]
toCopy Lambda SOACS
lam = do
  (copies, nameMap) <- [VName] -> m (Stms SOACS, Map VName VName)
forall (m :: * -> *).
(LocalScope SOACS m, MonadFreshNames m) =>
[VName] -> m (Stms SOACS, Map VName VName)
makeCopyStms [VName]
toCopy
  let l_body = Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
      newBody = Stms SOACS -> Body SOACS -> Body SOACS
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> GBody rep res -> GBody rep res
insertStms Stms SOACS
copies (Map VName VName -> Body SOACS -> Body SOACS
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
nameMap Body SOACS
l_body)
      newLambda = Lambda SOACS
lam {lambdaBody = newBody}
  pure newLambda

makeCopyStms ::
  (LocalScope SOACS m, MonadFreshNames m) =>
  [VName] ->
  m (Stms SOACS, M.Map VName VName)
makeCopyStms :: forall (m :: * -> *).
(LocalScope SOACS m, MonadFreshNames m) =>
[VName] -> m (Stms SOACS, Map VName VName)
makeCopyStms [VName]
vs = do
  vs' <- (VName -> m VName) -> [VName] -> m [VName]
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 -> m VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
makeNewName [VName]
vs
  copies <- forM (zip vs vs') $ \(VName
name, VName
name') ->
    [VName] -> Exp SOACS -> m (Stm SOACS)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m, HasScope rep m) =>
[VName] -> Exp rep -> m (Stm rep)
forall (m :: * -> *).
(MonadFreshNames m, HasScope SOACS m) =>
[VName] -> Exp SOACS -> m (Stm SOACS)
mkLetNames [VName
name'] (Exp SOACS -> m (Stm SOACS)) -> Exp SOACS -> m (Stm SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate ShapeBase SubExp
forall a. Monoid a => a
mempty (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
name
  pure (stmsFromList copies, M.fromList $ zip vs vs')
  where
    makeNewName :: VName -> m VName
makeNewName VName
name = Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> m VName) -> Name -> m VName
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName VName
name Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_copy"

okToFuseProducer :: H.SOAC SOACS -> FusionM Bool
okToFuseProducer :: SOAC SOACS -> FusionM Bool
okToFuseProducer (H.Screma SubExp
_ [Input]
_ ScremaForm SOACS
form) = do
  let is_scan :: Bool
is_scan = Maybe ([Scan SOACS], Lambda SOACS) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe ([Scan SOACS], Lambda SOACS) -> Bool)
-> Maybe ([Scan SOACS], Lambda SOACS) -> Bool
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> Maybe ([Scan SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
Futhark.isScanomapSOAC ScremaForm SOACS
form
  (FusionEnv -> Bool) -> FusionM Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((FusionEnv -> Bool) -> FusionM Bool)
-> (FusionEnv -> Bool) -> FusionM Bool
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool
not Bool
is_scan Bool -> Bool -> Bool
||) (Bool -> Bool) -> (FusionEnv -> Bool) -> FusionEnv -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FusionEnv -> Bool
fuseScans
okToFuseProducer SOAC SOACS
_ = Bool -> FusionM Bool
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

-- First node is producer, second is consumer.
vFuseNodeT ::
  [EdgeT] ->
  [VName] ->
  (NodeT, [EdgeT], [EdgeT]) ->
  (NodeT, [EdgeT]) ->
  FusionM (Maybe NodeT)
vFuseNodeT :: [EdgeT]
-> [VName]
-> (NodeT, [EdgeT], [EdgeT])
-> (NodeT, [EdgeT])
-> FusionM (Maybe NodeT)
vFuseNodeT [EdgeT]
_ [VName]
infusible (NodeT
s1, [EdgeT]
_, [EdgeT]
e1s) (MatchNode Stm SOACS
stm2 [(NodeT, [EdgeT])]
dfused, [EdgeT]
_)
  | NodeT -> Bool
isRealNode NodeT
s1,
    [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
infusible =
      Maybe NodeT -> FusionM (Maybe NodeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe NodeT -> FusionM (Maybe NodeT))
-> Maybe NodeT -> FusionM (Maybe NodeT)
forall a b. (a -> b) -> a -> b
$ NodeT -> Maybe NodeT
forall a. a -> Maybe a
Just (NodeT -> Maybe NodeT) -> NodeT -> Maybe NodeT
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> [(NodeT, [EdgeT])] -> NodeT
MatchNode Stm SOACS
stm2 ([(NodeT, [EdgeT])] -> NodeT) -> [(NodeT, [EdgeT])] -> NodeT
forall a b. (a -> b) -> a -> b
$ (NodeT
s1, [EdgeT]
e1s) (NodeT, [EdgeT]) -> [(NodeT, [EdgeT])] -> [(NodeT, [EdgeT])]
forall a. a -> [a] -> [a]
: [(NodeT, [EdgeT])]
dfused
vFuseNodeT [EdgeT]
_ [VName]
infusible (TransNode VName
stm1_out ArrayTransform
tr VName
stm1_in, [EdgeT]
_, [EdgeT]
_) (SoacNode ArrayTransforms
ots2 Pat Type
pats2 SOAC SOACS
soac2 StmAux (ExpDec SOACS)
aux2, [EdgeT]
_)
  | [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
infusible = do
      stm1_in_t <- VName -> FusionM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
stm1_in
      let onInput Input
inp
            | Input -> VName
H.inputArray Input
inp VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
stm1_out =
                ArrayTransforms -> VName -> Type -> Input
H.Input (ArrayTransform
tr ArrayTransform -> ArrayTransforms -> ArrayTransforms
H.<| Input -> ArrayTransforms
H.inputTransforms Input
inp) VName
stm1_in Type
stm1_in_t
            | Bool
otherwise =
                Input
inp
          soac2' = (Input -> Input) -> [Input] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map Input -> Input
onInput (SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
H.inputs SOAC SOACS
soac2) [Input] -> SOAC SOACS -> SOAC SOACS
forall rep. [Input] -> SOAC rep -> SOAC rep
`H.setInputs` SOAC SOACS
soac2
      pure $ Just $ SoacNode ots2 pats2 soac2' aux2
vFuseNodeT
  [EdgeT]
_
  [VName]
_
  (SoacNode ArrayTransforms
ots1 Pat Type
pats1 SOAC SOACS
soac1 StmAux (ExpDec SOACS)
aux1, [EdgeT]
i1s, [EdgeT]
_e1s)
  (SoacNode ArrayTransforms
ots2 Pat Type
pats2 SOAC SOACS
soac2 StmAux (ExpDec SOACS)
aux2, [EdgeT]
_e2s) = do
    let ker :: FusedSOAC
ker =
          TF.FusedSOAC
            { fsSOAC :: SOAC SOACS
TF.fsSOAC = SOAC SOACS
soac2,
              fsOutputTransform :: ArrayTransforms
TF.fsOutputTransform = ArrayTransforms
ots2,
              fsOutNames :: [VName]
TF.fsOutNames = Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pats2
            }
        preserveEdge :: EdgeT -> Bool
preserveEdge InfDep {} = Bool
True
        preserveEdge EdgeT
e = EdgeT -> Bool
isDep EdgeT
e
        preserve :: Names
preserve = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (EdgeT -> VName) -> [EdgeT] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map EdgeT -> VName
getName ([EdgeT] -> [VName]) -> [EdgeT] -> [VName]
forall a b. (a -> b) -> a -> b
$ (EdgeT -> Bool) -> [EdgeT] -> [EdgeT]
forall a. (a -> Bool) -> [a] -> [a]
filter EdgeT -> Bool
preserveEdge [EdgeT]
i1s
    ok <- SOAC SOACS -> FusionM Bool
okToFuseProducer SOAC SOACS
soac1
    -- It is not safe to fuse if any accumulators are updated by both, as the
    -- semantics require that any updates done by the consumer take precedence
    -- over those in the producer. This is implemented with a manual check here
    -- for convenience, but it could be argued that this should really be a Fake
    -- edge in the graph.
    let isProducedAcc (H.Input ArrayTransforms
_ VName
v Acc {}) = VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pats1
        isProducedAcc Input
_ = Bool
False
    r <-
      if ok && ots1 == mempty && not (any isProducedAcc (H.inputs soac2))
        then TF.attemptFusion TF.Vertical preserve (patNames pats1) soac1 ker
        else pure Nothing
    case r of
      Just FusedSOAC
ker' -> do
        let pats2' :: [PatElem Type]
pats2' =
              (VName -> Type -> PatElem Type)
-> [VName] -> [Type] -> [PatElem Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem (FusedSOAC -> [VName]
TF.fsOutNames FusedSOAC
ker') (SOAC SOACS -> [Type]
forall rep. SOAC rep -> [Type]
H.typeOf (FusedSOAC -> SOAC SOACS
TF.fsSOAC FusedSOAC
ker'))
        NodeT -> FusionM (Maybe NodeT)
fusedSomething (NodeT -> FusionM (Maybe NodeT)) -> NodeT -> FusionM (Maybe NodeT)
forall a b. (a -> b) -> a -> b
$
          ArrayTransforms
-> Pat Type -> SOAC SOACS -> StmAux (ExpDec SOACS) -> NodeT
SoacNode
            (FusedSOAC -> ArrayTransforms
TF.fsOutputTransform FusedSOAC
ker')
            ([PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem Type]
pats2')
            (FusedSOAC -> SOAC SOACS
TF.fsSOAC FusedSOAC
ker')
            (StmAux ()
StmAux (ExpDec SOACS)
aux1 StmAux () -> StmAux () -> StmAux ()
forall a. Semigroup a => a -> a -> a
<> StmAux ()
StmAux (ExpDec SOACS)
aux2)
      Maybe FusedSOAC
Nothing -> Maybe NodeT -> FusionM (Maybe NodeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NodeT
forall a. Maybe a
Nothing
vFuseNodeT
  [EdgeT]
_
  [VName]
infusible
  (SoacNode ArrayTransforms
ots1 Pat Type
pat1 (H.Screma SubExp
w [Input]
inps ScremaForm SOACS
form) StmAux (ExpDec SOACS)
aux1, [EdgeT]
_, [EdgeT]
_)
  (TransNode VName
stm2_out (H.Index StmAux ()
cs slice :: Slice SubExp
slice@(Slice (ds :: DimIndex SubExp
ds@(DimSlice SubExp
_ SubExp
w' SubExp
_) : [DimIndex SubExp]
ds_rest))) VName
_, [EdgeT]
_)
    | [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
infusible,
      SubExp
w SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
/= SubExp
w',
      ArrayTransforms
ots1 ArrayTransforms -> ArrayTransforms -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayTransforms
forall a. Monoid a => a
mempty,
      Just Lambda SOACS
_ <- ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form,
      [PatElem Type
pe] <- Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat1 = do
        let out_t :: Type
out_t = PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem Type
pe Type -> ShapeBase SubExp -> Type
forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
`setArrayShape` Slice SubExp -> ShapeBase SubExp
forall d. Slice d -> ShapeBase d
sliceShape Slice SubExp
slice
            inps' :: [Input]
inps' = (Input -> Input) -> [Input] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map Input -> Input
sliceInput [Input]
inps
            -- Even if we move the slice of the outermost dimension, there
            -- might still be some slicing of the inner ones.
            ots1' :: ArrayTransforms
ots1' = ArrayTransforms
ots1 ArrayTransforms -> ArrayTransform -> ArrayTransforms
H.|> StmAux () -> Slice SubExp -> ArrayTransform
H.Index StmAux ()
cs ([DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice (SubExp -> DimIndex SubExp
sliceDim SubExp
w' DimIndex SubExp -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. a -> [a] -> [a]
: [DimIndex SubExp]
ds_rest))
        NodeT -> FusionM (Maybe NodeT)
fusedSomething (NodeT -> FusionM (Maybe NodeT)) -> NodeT -> FusionM (Maybe NodeT)
forall a b. (a -> b) -> a -> b
$
          ArrayTransforms
-> Pat Type -> SOAC SOACS -> StmAux (ExpDec SOACS) -> NodeT
SoacNode
            ArrayTransforms
ots1'
            ([PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
stm2_out Type
out_t])
            (SubExp -> [Input] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [Input] -> ScremaForm rep -> SOAC rep
H.Screma SubExp
w' [Input]
inps' ScremaForm SOACS
form)
            StmAux (ExpDec SOACS)
aux1
    where
      sliceInput :: Input -> Input
sliceInput Input
inp =
        ArrayTransform -> Input -> Input
H.addTransform
          (StmAux () -> Slice SubExp -> ArrayTransform
H.Index StmAux ()
cs (Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice (Input -> Type
H.inputType Input
inp) [DimIndex SubExp
ds]))
          Input
inp
-- Case of fusing a screma with an WithAcc such as to (hopefully) perform
--   more fusion within the WithAcc. This would allow the withAcc to move in
--   the code (since up to now they mostly remain where they were introduced.)
-- We conservatively allow the fusion to fire---i.e., to move the soac inside
--   the withAcc---when the following are not part of withAcc's accumulators:
--    1. the in-dependencies of the soac and
--    2. the result of the soac
--  Note that the soac result is allowed to be part of the `infusible`
--    for as long as it is returned by the withAcc. If `infusible` is empty
--    then the extranous result will be simplified away.
vFuseNodeT
  [EdgeT]
edges
  [VName]
_infusible
  (SoacNode ArrayTransforms
ots1 Pat Type
pat1 soac :: SOAC SOACS
soac@(H.Screma SubExp
_w [Input]
_form ScremaForm SOACS
_s_inps) StmAux (ExpDec SOACS)
aux1, [EdgeT]
_is1, [EdgeT]
os1)
  (StmNode (Let Pat (LetDec SOACS)
pat2 StmAux (ExpDec SOACS)
aux2 (WithAcc [WithAccInput SOACS]
w_inps Lambda SOACS
lam0)), [EdgeT]
_os2)
    | ArrayTransforms
ots1 ArrayTransforms -> ArrayTransforms -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayTransforms
forall a. Monoid a => a
mempty,
      Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (EdgeT -> Bool) -> [EdgeT] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any EdgeT -> Bool
isFake [EdgeT]
edges,
      Names
wacc_cons_nms <- [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (WithAccInput SOACS -> [VName]) -> [WithAccInput SOACS] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(ShapeBase SubExp
_, [VName]
nms, Maybe (Lambda SOACS, [SubExp])
_) -> [VName]
nms) [WithAccInput SOACS]
w_inps,
      [VName]
soac_prod_nms <- (PatElem Type -> VName) -> [PatElem Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName ([PatElem Type] -> [VName]) -> [PatElem Type] -> [VName]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat1,
      [VName]
soac_indep_nms <- (EdgeT -> VName) -> [EdgeT] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map EdgeT -> VName
getName [EdgeT]
os1,
      (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Names -> Bool
`notNameIn` Names
wacc_cons_nms) ([VName]
soac_indep_nms [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
soac_prod_nms) = do
        lam <- (Lambda SOACS, Bool) -> Lambda SOACS
forall a b. (a, b) -> a
fst ((Lambda SOACS, Bool) -> Lambda SOACS)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
lam0
        bdy' <-
          runBodyBuilder $ inScopeOf lam $ do
            soac' <- H.toExp soac
            addStm $ Let pat1 aux1 soac'
            lam_res <- bodyBind $ lambdaBody lam
            let pat1_res = (VName -> SubExpRes) -> [VName] -> [SubExpRes]
forall a b. (a -> b) -> [a] -> [b]
map (Certs -> SubExp -> SubExpRes
SubExpRes ([VName] -> Certs
Certs []) (SubExp -> SubExpRes) -> (VName -> SubExp) -> VName -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) [VName]
soac_prod_nms
            pure $ lam_res ++ pat1_res
        let lam_ret_tp = Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ (PatElem Type -> Type) -> [PatElem Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat1)
            pat = [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat2 [PatElem Type] -> [PatElem Type] -> [PatElem Type]
forall a. [a] -> [a] -> [a]
++ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat1
        lam' <- renameLambda $ lam {lambdaBody = bdy', lambdaReturnType = lam_ret_tp}
        -- see if bringing the map inside the scatter has actually benefitted fusion
        (lam'', success) <- doFusionInLambda lam'
        if not success
          then pure Nothing
          else do
            -- `aux1` already appear in the moved SOAC stm; is there
            -- any need to add it to the enclosing withAcc stm as well?
            fusedSomething $ StmNode $ Let pat aux2 $ WithAcc w_inps lam''

--
-- The reverse of the case above, i.e., fusing a screma at the back of an
--   WithAcc such as to (hopefully) enable more fusion there.
-- This should be safe as long as the SOAC does not uses any of the
--   accumulator arrays produced by the withAcc.
-- We could not provide a test for this case, due to the very restrictive
--   way in which accumulators can be used at source level.
--
--
vFuseNodeT
  [EdgeT]
edges
  [VName]
_infusible
  (StmNode (Let Pat (LetDec SOACS)
pat1 StmAux (ExpDec SOACS)
aux1 (WithAcc [WithAccInput SOACS]
w_inps Lambda SOACS
wlam0)), [EdgeT]
_is1, [EdgeT]
_os1)
  (SoacNode ArrayTransforms
ots2 Pat Type
pat2 soac :: SOAC SOACS
soac@(H.Screma SubExp
_w [Input]
_form ScremaForm SOACS
_s_inps) StmAux (ExpDec SOACS)
aux2, [EdgeT]
_os2)
    | ArrayTransforms
ots2 ArrayTransforms -> ArrayTransforms -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayTransforms
forall a. Monoid a => a
mempty,
      Int
n <- [Param Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
wlam0) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2,
      Names
pat1_acc_nms <- [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
take Int
n ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ (PatElem Type -> VName) -> [PatElem Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName ([PatElem Type] -> [VName]) -> [PatElem Type] -> [VName]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat1,
      -- not $ namesIntersect (freeIn soac) pat1_acc_nms
      (EdgeT -> Bool) -> [EdgeT] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((VName -> Names -> Bool
`notNameIn` Names
pat1_acc_nms) (VName -> Bool) -> (EdgeT -> VName) -> EdgeT -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EdgeT -> VName
getName) [EdgeT]
edges = do
        wlam <- (Lambda SOACS, Bool) -> Lambda SOACS
forall a b. (a, b) -> a
fst ((Lambda SOACS, Bool) -> Lambda SOACS)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
wlam0
        bdy' <-
          runBodyBuilder $ inScopeOf wlam $ do
            -- adding stms of withacc's lambda
            wlam_res <- bodyBind $ lambdaBody wlam
            -- add copies of the non-accumulator results of withacc
            let other_pr1 = Int -> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
forall a. Int -> [a] -> [a]
drop Int
n ([(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)])
-> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
forall a b. (a -> b) -> a -> b
$ [PatElem Type] -> [SubExpRes] -> [(PatElem Type, SubExpRes)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat1) [SubExpRes]
wlam_res
            forM_ other_pr1 $ \(PatElem Type
pat_elm, SubExpRes
bdy_res) -> do
              let (VName
nm, SubExp
se, Type
tp) = (PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pat_elm, SubExpRes -> SubExp
resSubExp SubExpRes
bdy_res, PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem Type
pat_elm)
                  aux :: StmAux ()
aux = (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) {stmAuxCerts = resCerts bdy_res}
              Stm (Rep (BuilderT SOACS (State VNameSource))) -> Builder SOACS ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm (Rep (BuilderT SOACS (State VNameSource)))
 -> Builder SOACS ())
-> Stm (Rep (BuilderT SOACS (State VNameSource)))
-> Builder SOACS ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
-> StmAux (ExpDec (Rep (BuilderT SOACS (State VNameSource))))
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> Stm (Rep (BuilderT SOACS (State VNameSource)))
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
nm Type
tp]) StmAux ()
StmAux (ExpDec (Rep (BuilderT SOACS (State VNameSource))))
aux (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> Stm (Rep (BuilderT SOACS (State VNameSource))))
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> Stm (Rep (BuilderT SOACS (State VNameSource)))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
-> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
            -- add the soac stmt
            soac' <- H.toExp soac
            addStm $ Let pat2 aux2 soac'
            -- build the body result
            let pat2_res = (PatElem Type -> SubExpRes) -> [PatElem Type] -> [SubExpRes]
forall a b. (a -> b) -> [a] -> [b]
map (Certs -> SubExp -> SubExpRes
SubExpRes ([VName] -> Certs
Certs []) (SubExp -> SubExpRes)
-> (PatElem Type -> SubExp) -> PatElem Type -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var (VName -> SubExp)
-> (PatElem Type -> VName) -> PatElem Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName) ([PatElem Type] -> [SubExpRes]) -> [PatElem Type] -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat2
            pure $ wlam_res ++ pat2_res
        let lam_ret_tp = Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
wlam [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ (PatElem Type -> Type) -> [PatElem Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat2)
            pat = [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat1 [PatElem Type] -> [PatElem Type] -> [PatElem Type]
forall a. [a] -> [a] -> [a]
++ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat2
        wlam' <- renameLambda $ wlam {lambdaBody = bdy', lambdaReturnType = lam_ret_tp}
        -- see if bringing the map inside the scatter has actually benefitted fusion
        (wlam'', success) <- doFusionInLambda wlam'
        if not success
          then pure Nothing
          else -- `aux2` already appear in the enclosed SOAC stm; is there
          -- any need to add it to the enclosing withAcc stm as well?
            fusedSomething $ StmNode $ Let pat aux1 $ WithAcc w_inps wlam''
-- the case of fusing two withaccs
vFuseNodeT
  [EdgeT]
edges
  [VName]
infusible
  (StmNode (Let Pat (LetDec SOACS)
pat1 StmAux (ExpDec SOACS)
aux1 (WithAcc [WithAccInput SOACS]
w1_inps Lambda SOACS
lam1)), [EdgeT]
is1, [EdgeT]
_os1)
  (StmNode (Let Pat (LetDec SOACS)
pat2 StmAux (ExpDec SOACS)
aux2 (WithAcc [WithAccInput SOACS]
w2_inps Lambda SOACS
lam2)), [EdgeT]
_os2)
    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (EdgeT -> Bool) -> [EdgeT] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any EdgeT -> Bool
isFake [EdgeT]
edges,
      Names
wacc2_cons_nms <- [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (WithAccInput SOACS -> [VName]) -> [WithAccInput SOACS] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(ShapeBase SubExp
_, [VName]
nms, Maybe (Lambda SOACS, [SubExp])
_) -> [VName]
nms) [WithAccInput SOACS]
w2_inps,
      [VName]
wacc1_indep_nms <- (EdgeT -> VName) -> [EdgeT] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map EdgeT -> VName
getName [EdgeT]
is1,
      (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Names -> Bool
`notNameIn` Names
wacc2_cons_nms) [VName]
wacc1_indep_nms = do
        -- the other safety checks are done inside `tryFuseWithAccs`
        lam1' <- (Lambda SOACS, Bool) -> Lambda SOACS
forall a b. (a, b) -> a
fst ((Lambda SOACS, Bool) -> Lambda SOACS)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
lam1
        lam2' <- fst <$> doFusionInLambda lam2
        let stm1 = Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat1 StmAux (ExpDec SOACS)
aux1 ([WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [WithAccInput SOACS]
w1_inps Lambda SOACS
lam1')
            stm2 = Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat2 StmAux (ExpDec SOACS)
aux2 ([WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [WithAccInput SOACS]
w2_inps Lambda SOACS
lam2')
        mstm <- sequence $ SF.tryFuseWithAccs infusible stm1 stm2
        case mstm of
          Maybe (Stm SOACS)
Nothing -> Maybe NodeT -> FusionM (Maybe NodeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NodeT
forall a. Maybe a
Nothing
          Just (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (WithAcc [WithAccInput SOACS]
w_inps Lambda SOACS
wlam)) -> do
            (wlam', success) <- Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
wlam
            let new_stm = Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux ([WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [WithAccInput SOACS]
w_inps Lambda SOACS
wlam')
            if success then fusedSomething (StmNode new_stm) else pure Nothing
          Just Stm SOACS
_ -> String -> FusionM (Maybe NodeT)
forall a. HasCallStack => String -> a
error String
"Illegal result of tryFuseWithAccs called from vFuseNodeT."
--
vFuseNodeT [EdgeT]
_ [VName]
_ (NodeT, [EdgeT], [EdgeT])
_ (NodeT, [EdgeT])
_ = Maybe NodeT -> FusionM (Maybe NodeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NodeT
forall a. Maybe a
Nothing

resFromLambda :: Lambda rep -> Result
resFromLambda :: forall rep. Lambda rep -> [SubExpRes]
resFromLambda = GBody rep SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody rep SubExpRes -> [SubExpRes])
-> (Lambda rep -> GBody rep SubExpRes) -> Lambda rep -> [SubExpRes]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody

hasNoDifferingInputs :: [H.Input] -> [H.Input] -> Bool
hasNoDifferingInputs :: [Input] -> [Input] -> Bool
hasNoDifferingInputs [Input]
is1 [Input]
is2 =
  let ([Input]
vs1, [Input]
vs2) = ([Input] -> [Input]
isNotVarInput [Input]
is1, [Input] -> [Input]
isNotVarInput ([Input] -> [Input]) -> [Input] -> [Input]
forall a b. (a -> b) -> a -> b
$ [Input]
is2 [Input] -> [Input] -> [Input]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ [Input]
is1)
   in [Input] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Input] -> Bool) -> [Input] -> Bool
forall a b. (a -> b) -> a -> b
$ [Input]
vs1 [Input] -> [Input] -> [Input]
forall a. Eq a => [a] -> [a] -> [a]
`L.intersect` [Input]
vs2

hFuseNodeT :: NodeT -> NodeT -> FusionM (Maybe NodeT)
hFuseNodeT :: NodeT -> NodeT -> FusionM (Maybe NodeT)
hFuseNodeT (SoacNode ArrayTransforms
ots1 Pat Type
pats1 SOAC SOACS
soac1 StmAux (ExpDec SOACS)
aux1) (SoacNode ArrayTransforms
ots2 Pat Type
pats2 SOAC SOACS
soac2 StmAux (ExpDec SOACS)
aux2)
  | ArrayTransforms
ots1 ArrayTransforms -> ArrayTransforms -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayTransforms
forall a. Monoid a => a
mempty,
    ArrayTransforms
ots2 ArrayTransforms -> ArrayTransforms -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayTransforms
forall a. Monoid a => a
mempty,
    [Input] -> [Input] -> Bool
hasNoDifferingInputs (SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
H.inputs SOAC SOACS
soac1) (SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
H.inputs SOAC SOACS
soac2) = do
      let ker :: FusedSOAC
ker =
            TF.FusedSOAC
              { fsSOAC :: SOAC SOACS
TF.fsSOAC = SOAC SOACS
soac2,
                fsOutputTransform :: ArrayTransforms
TF.fsOutputTransform = ArrayTransforms
forall a. Monoid a => a
mempty,
                fsOutNames :: [VName]
TF.fsOutNames = Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pats2
              }
          preserve :: Names
preserve = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pats1
      r <- Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> FusionM (Maybe FusedSOAC)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> m (Maybe FusedSOAC)
TF.attemptFusion Mode
TF.Horizontal Names
preserve (Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pats1) SOAC SOACS
soac1 FusedSOAC
ker
      case r of
        Just FusedSOAC
ker' -> do
          let pats2' :: [PatElem Type]
pats2' =
                (VName -> Type -> PatElem Type)
-> [VName] -> [Type] -> [PatElem Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem (FusedSOAC -> [VName]
TF.fsOutNames FusedSOAC
ker') (SOAC SOACS -> [Type]
forall rep. SOAC rep -> [Type]
H.typeOf (FusedSOAC -> SOAC SOACS
TF.fsSOAC FusedSOAC
ker'))
          NodeT -> FusionM (Maybe NodeT)
fusedSomething (NodeT -> FusionM (Maybe NodeT)) -> NodeT -> FusionM (Maybe NodeT)
forall a b. (a -> b) -> a -> b
$ ArrayTransforms
-> Pat Type -> SOAC SOACS -> StmAux (ExpDec SOACS) -> NodeT
SoacNode ArrayTransforms
forall a. Monoid a => a
mempty ([PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem Type]
pats2') (FusedSOAC -> SOAC SOACS
TF.fsSOAC FusedSOAC
ker') (StmAux ()
StmAux (ExpDec SOACS)
aux1 StmAux () -> StmAux () -> StmAux ()
forall a. Semigroup a => a -> a -> a
<> StmAux ()
StmAux (ExpDec SOACS)
aux2)
        Maybe FusedSOAC
Nothing -> Maybe NodeT -> FusionM (Maybe NodeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NodeT
forall a. Maybe a
Nothing
hFuseNodeT
  (StmNode (Let Pat (LetDec SOACS)
pat1 StmAux (ExpDec SOACS)
aux1 (WithAcc [WithAccInput SOACS]
w1_inps Lambda SOACS
lam1)))
  (StmNode (Let Pat (LetDec SOACS)
pat2 StmAux (ExpDec SOACS)
aux2 (WithAcc [WithAccInput SOACS]
w2_inps Lambda SOACS
lam2))) = do
    -- The only tricky thing here is that we have to put all the
    -- accumulator-based results first.
    let num_inputs1 :: Int
num_inputs1 = [WithAccInput SOACS] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput SOACS]
w1_inps
        num_inputs2 :: Int
num_inputs2 = [WithAccInput SOACS] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput SOACS]
w2_inps
        num_arrs1 :: Int
num_arrs1 = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (WithAccInput SOACS -> Int) -> [WithAccInput SOACS] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (\(ShapeBase SubExp
_, [VName]
as, Maybe (Lambda SOACS, [SubExp])
_) -> [VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
as) [WithAccInput SOACS]
w1_inps
        num_arrs2 :: Int
num_arrs2 = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (WithAccInput SOACS -> Int) -> [WithAccInput SOACS] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (\(ShapeBase SubExp
_, [VName]
as, Maybe (Lambda SOACS, [SubExp])
_) -> [VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
as) [WithAccInput SOACS]
w2_inps
        w3_inps :: [WithAccInput SOACS]
w3_inps = [WithAccInput SOACS]
w1_inps [WithAccInput SOACS]
-> [WithAccInput SOACS] -> [WithAccInput SOACS]
forall a. Semigroup a => a -> a -> a
<> [WithAccInput SOACS]
w2_inps
        reorder :: (p -> [a]) -> Int -> p -> Int -> p -> [a]
reorder p -> [a]
f Int
n p
a Int
m p
b =
          let ([a]
a_xs, [a]
a_ys) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n ([a] -> ([a], [a])) -> [a] -> ([a], [a])
forall a b. (a -> b) -> a -> b
$ p -> [a]
f p
a
              ([a]
b_xs, [a]
b_ys) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
m ([a] -> ([a], [a])) -> [a] -> ([a], [a])
forall a b. (a -> b) -> a -> b
$ p -> [a]
f p
b
           in [a]
a_xs [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a]
b_xs [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a]
a_ys [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a]
b_ys
        lam3 :: Lambda SOACS
lam3 =
          [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda
            ((Lambda SOACS -> [Param Type])
-> Int -> Lambda SOACS -> Int -> Lambda SOACS -> [Param Type]
forall {p} {a}. (p -> [a]) -> Int -> p -> Int -> p -> [a]
reorder Lambda SOACS -> [Param Type]
Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Int
num_inputs1 Lambda SOACS
lam1 Int
num_inputs2 Lambda SOACS
lam2)
            ((Lambda SOACS -> [Type])
-> Int -> Lambda SOACS -> Int -> Lambda SOACS -> [Type]
forall {p} {a}. (p -> [a]) -> Int -> p -> Int -> p -> [a]
reorder Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Int
num_inputs1 Lambda SOACS
lam1 Int
num_inputs2 Lambda SOACS
lam2)
            (Body SOACS -> Lambda SOACS) -> Body SOACS -> Lambda SOACS
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody
              (Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam1) Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam2))
              ((Lambda SOACS -> [SubExpRes])
-> Int -> Lambda SOACS -> Int -> Lambda SOACS -> [SubExpRes]
forall {p} {a}. (p -> [a]) -> Int -> p -> Int -> p -> [a]
reorder (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes])
-> (Lambda SOACS -> Body SOACS) -> Lambda SOACS -> [SubExpRes]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody) Int
num_inputs1 Lambda SOACS
lam1 Int
num_inputs2 Lambda SOACS
lam2)
    NodeT -> FusionM (Maybe NodeT)
fusedSomething (NodeT -> FusionM (Maybe NodeT)) -> NodeT -> FusionM (Maybe NodeT)
forall a b. (a -> b) -> a -> b
$
      Stm SOACS -> NodeT
StmNode (Stm SOACS -> NodeT) -> Stm SOACS -> NodeT
forall a b. (a -> b) -> a -> b
$
        Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElem (LetDec SOACS)] -> Pat (LetDec SOACS)
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem (LetDec SOACS)] -> Pat (LetDec SOACS))
-> [PatElem (LetDec SOACS)] -> Pat (LetDec SOACS)
forall a b. (a -> b) -> a -> b
$ (Pat Type -> [PatElem (LetDec SOACS)])
-> Int -> Pat Type -> Int -> Pat Type -> [PatElem (LetDec SOACS)]
forall {p} {a}. (p -> [a]) -> Int -> p -> Int -> p -> [a]
reorder Pat Type -> [PatElem Type]
Pat Type -> [PatElem (LetDec SOACS)]
forall dec. Pat dec -> [PatElem dec]
patElems Int
num_arrs1 Pat Type
Pat (LetDec SOACS)
pat1 Int
num_arrs2 Pat Type
Pat (LetDec SOACS)
pat2) (StmAux ()
StmAux (ExpDec SOACS)
aux1 StmAux () -> StmAux () -> StmAux ()
forall a. Semigroup a => a -> a -> a
<> StmAux ()
StmAux (ExpDec SOACS)
aux2) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
          [WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [WithAccInput SOACS]
w3_inps Lambda SOACS
lam3
hFuseNodeT NodeT
_ NodeT
_ = Maybe NodeT -> FusionM (Maybe NodeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NodeT
forall a. Maybe a
Nothing

removeOutputsExcept :: [VName] -> NodeT -> FusionM NodeT
removeOutputsExcept :: [VName] -> NodeT -> FusionM NodeT
removeOutputsExcept [VName]
toKeep NodeT
s = case NodeT
s of
  SoacNode ArrayTransforms
ots (Pat [PatElem Type]
pats) (H.Screma SubExp
w [Input]
inp (ScremaForm Lambda SOACS
pre_lam [] [Reduce SOACS]
red Lambda SOACS
post_lam)) StmAux (ExpDec SOACS)
aux1 -> do
    pre_lam' <- if Bool
changed then Lambda SOACS -> FusionM (Lambda SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Lambda SOACS -> m (Lambda SOACS)
simplifyLambda Lambda SOACS
new_pre else Lambda SOACS -> FusionM (Lambda SOACS)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda SOACS
new_pre
    post_lam' <- if changed then simplifyLambda new_post else pure new_post
    pure $
      SoacNode
        ots
        (Pat $ red_pats <> new_pats)
        (H.Screma w inp (ScremaForm pre_lam' [] red post_lam'))
        aux1
    where
      ([SubExpRes]
pre_red_res, [SubExpRes]
pre_map_res) =
        Int -> [SubExpRes] -> ([SubExpRes], [SubExpRes])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
red) ([SubExpRes] -> ([SubExpRes], [SubExpRes]))
-> [SubExpRes] -> ([SubExpRes], [SubExpRes])
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [SubExpRes]
forall rep. Lambda rep -> [SubExpRes]
resFromLambda Lambda SOACS
pre_lam
      ([Type]
pre_red_ts, [Type]
pre_map_ts) =
        Int -> [Type] -> ([Type], [Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
red) ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
pre_lam

      to_change :: [(SubExpRes, Type, Param Type, SubExpRes, Type)]
to_change =
        [SubExpRes]
-> [Type]
-> [Param Type]
-> [SubExpRes]
-> [Type]
-> [(SubExpRes, Type, Param Type, SubExpRes, Type)]
forall a b c d e.
[a] -> [b] -> [c] -> [d] -> [e] -> [(a, b, c, d, e)]
L.zip5
          [SubExpRes]
pre_map_res
          [Type]
pre_map_ts
          (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
post_lam)
          (Lambda SOACS -> [SubExpRes]
forall rep. Lambda rep -> [SubExpRes]
resFromLambda Lambda SOACS
post_lam)
          (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
post_lam)

      ([PatElem Type]
red_pats, [PatElem Type]
map_pats) = Int -> [PatElem Type] -> ([PatElem Type], [PatElem Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
red) [PatElem Type]
pats

      changed :: Bool
changed = Lambda SOACS
new_post Lambda SOACS -> Lambda SOACS -> Bool
forall a. Eq a => a -> a -> Bool
/= Lambda SOACS
post_lam Bool -> Bool -> Bool
|| Lambda SOACS
new_pre Lambda SOACS -> Lambda SOACS -> Bool
forall a. Eq a => a -> a -> Bool
/= Lambda SOACS
pre_lam

      ([PatElem Type]
new_pats, [(SubExpRes, Type, Param Type, SubExpRes, Type)]
new) =
        [(PatElem Type, (SubExpRes, Type, Param Type, SubExpRes, Type))]
-> ([PatElem Type],
    [(SubExpRes, Type, Param Type, SubExpRes, Type)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PatElem Type, (SubExpRes, Type, Param Type, SubExpRes, Type))]
 -> ([PatElem Type],
     [(SubExpRes, Type, Param Type, SubExpRes, Type)]))
-> [(PatElem Type, (SubExpRes, Type, Param Type, SubExpRes, Type))]
-> ([PatElem Type],
    [(SubExpRes, Type, Param Type, SubExpRes, Type)])
forall a b. (a -> b) -> a -> b
$
          ((PatElem Type, (SubExpRes, Type, Param Type, SubExpRes, Type))
 -> Bool)
-> [(PatElem Type, (SubExpRes, Type, Param Type, SubExpRes, Type))]
-> [(PatElem Type, (SubExpRes, Type, Param Type, SubExpRes, Type))]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(PatElem Type
x, (SubExpRes, Type, Param Type, SubExpRes, Type)
_) -> PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
x VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
toKeep) ([PatElem Type]
-> [(SubExpRes, Type, Param Type, SubExpRes, Type)]
-> [(PatElem Type, (SubExpRes, Type, Param Type, SubExpRes, Type))]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatElem Type]
map_pats [(SubExpRes, Type, Param Type, SubExpRes, Type)]
to_change)
      ( [SubExpRes]
new_pre_map_res,
        [Type]
new_pre_map_ts,
        [Param Type]
new_post_pars,
        [SubExpRes]
new_post_res,
        [Type]
new_post_ts
        ) = [(SubExpRes, Type, Param Type, SubExpRes, Type)]
-> ([SubExpRes], [Type], [Param Type], [SubExpRes], [Type])
forall a b c d e. [(a, b, c, d, e)] -> ([a], [b], [c], [d], [e])
L.unzip5 [(SubExpRes, Type, Param Type, SubExpRes, Type)]
new
      new_post :: Lambda SOACS
new_post =
        Lambda
          { lambdaParams :: [LParam SOACS]
lambdaParams = [Param Type]
[LParam SOACS]
new_post_pars,
            lambdaReturnType :: [Type]
lambdaReturnType = [Type]
new_post_ts,
            lambdaBody :: Body SOACS
lambdaBody = (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
post_lam) {bodyResult = new_post_res}
          }
      new_pre :: Lambda SOACS
new_pre =
        Lambda SOACS
pre_lam
          { lambdaReturnType = pre_red_ts <> new_pre_map_ts,
            lambdaBody =
              (lambdaBody pre_lam)
                { bodyResult = pre_red_res <> new_pre_map_res
                }
          }
  SoacNode ArrayTransforms
ots (Pat [PatElem Type]
pats1) (H.Screma SubExp
w [Input]
inp (ScremaForm Lambda SOACS
pre_lam [Scan SOACS]
scan [Reduce SOACS]
red Lambda SOACS
post_lam)) StmAux (ExpDec SOACS)
aux1 -> do
    post_lam' <- if Bool
changed then Lambda SOACS -> FusionM (Lambda SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Lambda SOACS -> m (Lambda SOACS)
simplifyLambda Lambda SOACS
new_post else Lambda SOACS -> FusionM (Lambda SOACS)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda SOACS
new_post
    pure $
      SoacNode
        ots
        (Pat $ pats_unchanged <> pats_new)
        (H.Screma w inp (ScremaForm pre_lam scan red post_lam'))
        aux1
    where
      red_output_size :: Int
red_output_size = [Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
Futhark.redResults [Reduce SOACS]
red

      ([PatElem Type]
pats_unchanged, [PatElem Type]
pats_toChange) = Int -> [PatElem Type] -> ([PatElem Type], [PatElem Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
red_output_size [PatElem Type]
pats1
      res_toChange :: [(SubExpRes, Type)]
res_toChange = [SubExpRes] -> [Type] -> [(SubExpRes, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda SOACS -> [SubExpRes]
forall rep. Lambda rep -> [SubExpRes]
resFromLambda Lambda SOACS
post_lam) (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
post_lam)

      changed :: Bool
changed = Lambda SOACS
new_post Lambda SOACS -> Lambda SOACS -> Bool
forall a. Eq a => a -> a -> Bool
/= Lambda SOACS
post_lam

      ([PatElem Type]
pats_new, [(SubExpRes, Type)]
res_new) =
        [(PatElem Type, (SubExpRes, Type))]
-> ([PatElem Type], [(SubExpRes, Type)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PatElem Type, (SubExpRes, Type))]
 -> ([PatElem Type], [(SubExpRes, Type)]))
-> [(PatElem Type, (SubExpRes, Type))]
-> ([PatElem Type], [(SubExpRes, Type)])
forall a b. (a -> b) -> a -> b
$ ((PatElem Type, (SubExpRes, Type)) -> Bool)
-> [(PatElem Type, (SubExpRes, Type))]
-> [(PatElem Type, (SubExpRes, Type))]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(PatElem Type
x, (SubExpRes, Type)
_) -> PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
x VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
toKeep) ([PatElem Type]
-> [(SubExpRes, Type)] -> [(PatElem Type, (SubExpRes, Type))]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatElem Type]
pats_toChange [(SubExpRes, Type)]
res_toChange)
      ([SubExpRes]
results, [Type]
types) = [(SubExpRes, Type)] -> ([SubExpRes], [Type])
forall a b. [(a, b)] -> ([a], [b])
unzip [(SubExpRes, Type)]
res_new
      new_post :: Lambda SOACS
new_post =
        Lambda SOACS
post_lam
          { lambdaReturnType = types,
            lambdaBody = (lambdaBody post_lam) {bodyResult = results}
          }
  NodeT
node -> NodeT -> FusionM NodeT
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodeT
node

vNameFromAdj :: G.Node -> (EdgeT, G.Node) -> VName
vNameFromAdj :: Int -> (EdgeT, Int) -> VName
vNameFromAdj Int
n1 (EdgeT
edge, Int
n2) = DepEdge -> VName
depsFromEdge (Int
n2, Int
n1, EdgeT
edge)

removeUnusedOutputs :: DepGraphAug FusionM
removeUnusedOutputs :: DepGraphAug FusionM
removeUnusedOutputs = (Context NodeT EdgeT -> FusionM (Context NodeT EdgeT))
-> DepGraphAug FusionM
forall (m :: * -> *).
Monad m =>
(Context NodeT EdgeT -> m (Context NodeT EdgeT)) -> DepGraphAug m
mapAcross ((Context NodeT EdgeT -> FusionM (Context NodeT EdgeT))
 -> DepGraphAug FusionM)
-> (Context NodeT EdgeT -> FusionM (Context NodeT EdgeT))
-> DepGraphAug FusionM
forall a b. (a -> b) -> a -> b
$ \(Adj EdgeT
incoming, Int
n1, NodeT
nodeT, Adj EdgeT
outgoing) -> do
  let toKeep :: [VName]
toKeep = ((EdgeT, Int) -> VName) -> Adj EdgeT -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> (EdgeT, Int) -> VName
vNameFromAdj Int
n1) Adj EdgeT
incoming
  nodeT' <- [VName] -> NodeT -> FusionM NodeT
removeOutputsExcept [VName]
toKeep NodeT
nodeT
  pure (incoming, n1, nodeT', outgoing)

tryFuseNodeInGraph :: DepNode -> DepGraphAug FusionM
tryFuseNodeInGraph :: DepNode -> DepGraphAug FusionM
tryFuseNodeInGraph DepNode
node_to_fuse dg :: DepGraph
dg@DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
dgGraph = Gr NodeT EdgeT
g}
  | Bool -> Bool
not (Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem (DepNode -> Int
nodeFromLNode DepNode
node_to_fuse) Gr NodeT EdgeT
g) = DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
-- \^ Node might have been fused away since.
tryFuseNodeInGraph DepNode
node_to_fuse dg :: DepGraph
dg@DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
dgGraph = Gr NodeT EdgeT
g} =
  [DepGraphAug FusionM] -> DepGraphAug FusionM
forall (m :: * -> *). Monad m => [DepGraphAug m] -> DepGraphAug m
applyAugs ((Int -> DepGraphAug FusionM) -> [Int] -> [DepGraphAug FusionM]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> DepGraphAug FusionM
vTryFuseNodesInGraph Int
node_to_fuse_id) [Int]
fuses_with) DepGraph
dg
  where
    node_to_fuse_id :: Int
node_to_fuse_id = DepNode -> Int
nodeFromLNode DepNode
node_to_fuse
    relevant :: (Int, EdgeT) -> Bool
relevant (Int
n, InfDep VName
_) = Int -> DepGraph -> Bool
isWithAccNodeId Int
n DepGraph
dg
    relevant (Int
_, EdgeT
e) = EdgeT -> Bool
isDep EdgeT
e
    fuses_with :: [Int]
fuses_with = ((Int, EdgeT) -> Int) -> [(Int, EdgeT)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, EdgeT) -> Int
forall a b. (a, b) -> a
fst ([(Int, EdgeT)] -> [Int]) -> [(Int, EdgeT)] -> [Int]
forall a b. (a -> b) -> a -> b
$ ((Int, EdgeT) -> Bool) -> [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int, EdgeT) -> Bool
relevant ([(Int, EdgeT)] -> [(Int, EdgeT)])
-> [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a b. (a -> b) -> a -> b
$ Gr NodeT EdgeT -> Int -> [(Int, EdgeT)]
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> [(Int, b)]
G.lpre Gr NodeT EdgeT
g Int
node_to_fuse_id

doSoacThroughTransFusion :: DepGraphAug FusionM
doSoacThroughTransFusion :: DepGraphAug FusionM
doSoacThroughTransFusion DepGraph
dg =
  [DepGraphAug FusionM] -> DepGraphAug FusionM
forall (m :: * -> *). Monad m => [DepGraphAug m] -> DepGraphAug m
applyAugs
    [ (Lambda SOACS -> FusionM (Lambda SOACS, Bool))
-> (NodeT -> FusionM (Maybe NodeT)) -> Int -> DepGraphAug FusionM
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
(Lambda SOACS -> m (Lambda SOACS, Bool))
-> (NodeT -> m (Maybe NodeT)) -> Int -> DepGraphAug m
SF.trySoacThroughTransIntoWithAcc Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda NodeT -> FusionM (Maybe NodeT)
fusedSomething Int
wacc_id
    | (Int
wacc_id, StmNode (Let Pat (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ (WithAcc {}))) <- Gr NodeT EdgeT -> [DepNode]
forall a b. Gr a b -> [LNode a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LNode a]
G.labNodes (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg)
    ]
    DepGraph
dg

doVerticalFusion :: DepGraphAug FusionM
doVerticalFusion :: DepGraphAug FusionM
doVerticalFusion DepGraph
dg = [DepGraphAug FusionM] -> DepGraphAug FusionM
forall (m :: * -> *). Monad m => [DepGraphAug m] -> DepGraphAug m
applyAugs ((DepNode -> DepGraphAug FusionM)
-> [DepNode] -> [DepGraphAug FusionM]
forall a b. (a -> b) -> [a] -> [b]
map DepNode -> DepGraphAug FusionM
tryFuseNodeInGraph ([DepNode] -> [DepGraphAug FusionM])
-> [DepNode] -> [DepGraphAug FusionM]
forall a b. (a -> b) -> a -> b
$ [DepNode] -> [DepNode]
forall a. [a] -> [a]
reverse ([DepNode] -> [DepNode]) -> [DepNode] -> [DepNode]
forall a b. (a -> b) -> a -> b
$ (DepNode -> Bool) -> [DepNode] -> [DepNode]
forall a. (a -> Bool) -> [a] -> [a]
filter DepNode -> Bool
forall {a}. (a, NodeT) -> Bool
relevant ([DepNode] -> [DepNode]) -> [DepNode] -> [DepNode]
forall a b. (a -> b) -> a -> b
$ Gr NodeT EdgeT -> [DepNode]
forall a b. Gr a b -> [LNode a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LNode a]
G.labNodes (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg)) DepGraph
dg
  where
    relevant :: (a, NodeT) -> Bool
relevant (a
_, n :: NodeT
n@(StmNode {})) = NodeT -> Bool
isWithAccNodeT NodeT
n
    relevant (a
_, ResNode {}) = Bool
False
    relevant (a, NodeT)
_ = Bool
True

-- | For each pair of SOAC nodes that share an input, or any WithAcc nodes,
-- attempt to fuse them horizontally.
doHorizontalFusion :: DepGraphAug FusionM
doHorizontalFusion :: DepGraphAug FusionM
doHorizontalFusion DepGraph
dg = [DepGraphAug FusionM] -> DepGraphAug FusionM
forall (m :: * -> *). Monad m => [DepGraphAug m] -> DepGraphAug m
applyAugs ([DepGraphAug FusionM]
soac_pairs [DepGraphAug FusionM]
-> [DepGraphAug FusionM] -> [DepGraphAug FusionM]
forall a. Semigroup a => a -> a -> a
<> [DepGraphAug FusionM]
withacc_pairs) DepGraph
dg
  where
    soac_pairs, withacc_pairs :: [DepGraphAug FusionM]
    soac_pairs :: [DepGraphAug FusionM]
soac_pairs = do
      (x, SoacNode _ _ soac_x _) <- Gr NodeT EdgeT -> [DepNode]
forall a b. Gr a b -> [LNode a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LNode a]
G.labNodes (Gr NodeT EdgeT -> [DepNode]) -> Gr NodeT EdgeT -> [DepNode]
forall a b. (a -> b) -> a -> b
$ DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg
      (y, SoacNode _ _ soac_y _) <- G.labNodes $ dgGraph dg
      guard $ x < y
      -- Must share an input.
      guard $
        any
          ((`elem` map H.inputArray (H.inputs soac_x)) . H.inputArray)
          (H.inputs soac_y)
      pure $ \DepGraph
dg' -> do
        -- Nodes might have been fused away by now.
        if Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
x (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg') Bool -> Bool -> Bool
&& Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
y (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg')
          then Int -> Int -> DepGraphAug FusionM
hTryFuseNodesInGraph Int
x Int
y DepGraph
dg'
          else DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg'

    withacc_pairs :: [DepGraphAug FusionM]
withacc_pairs = do
      (x, StmNode (Let _ _ (WithAcc {}))) <- Gr NodeT EdgeT -> [DepNode]
forall a b. Gr a b -> [LNode a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LNode a]
G.labNodes (Gr NodeT EdgeT -> [DepNode]) -> Gr NodeT EdgeT -> [DepNode]
forall a b. (a -> b) -> a -> b
$ DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg
      (y, StmNode (Let _ _ (WithAcc {}))) <- G.labNodes $ dgGraph dg
      guard $ x < y
      pure $ \DepGraph
dg' -> do
        -- Nodes might have been fused away by now.
        if Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
x (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg') Bool -> Bool -> Bool
&& Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
G.gelem Int
y (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg')
          then Int -> Int -> DepGraphAug FusionM
hTryFuseNodesInGraph Int
x Int
y DepGraph
dg'
          else DepGraphAug FusionM
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg'

doInnerFusion :: DepGraphAug FusionM
doInnerFusion :: DepGraphAug FusionM
doInnerFusion = (Context NodeT EdgeT -> FusionM (Context NodeT EdgeT))
-> DepGraphAug FusionM
forall (m :: * -> *).
Monad m =>
(Context NodeT EdgeT -> m (Context NodeT EdgeT)) -> DepGraphAug m
mapAcross Context NodeT EdgeT -> FusionM (Context NodeT EdgeT)
runInnerFusionOnContext

-- Fixed-point iteration.
keepTrying :: DepGraphAug FusionM -> DepGraphAug FusionM
keepTrying :: DepGraphAug FusionM -> DepGraphAug FusionM
keepTrying DepGraphAug FusionM
f DepGraph
g = do
  DepGraph -> FusionM DepGraph -> FusionM DepGraph
forall a. a -> FusionM a -> FusionM a
useGas DepGraph
g (FusionM DepGraph -> FusionM DepGraph)
-> FusionM DepGraph -> FusionM DepGraph
forall a b. (a -> b) -> a -> b
$ do
    prev_fused <- (FusionEnv -> Int) -> FusionM Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets FusionEnv -> Int
fusionCount
    g' <- f g
    aft_fused <- gets fusionCount
    if prev_fused /= aft_fused
      then keepTrying f g'
      else pure g'

doAllFusion :: DepGraphAug FusionM
doAllFusion :: DepGraphAug FusionM
doAllFusion =
  DepGraphAug FusionM -> DepGraphAug FusionM
keepTrying (DepGraphAug FusionM -> DepGraphAug FusionM)
-> ([DepGraphAug FusionM] -> DepGraphAug FusionM)
-> [DepGraphAug FusionM]
-> DepGraphAug FusionM
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [DepGraphAug FusionM] -> DepGraphAug FusionM
forall (m :: * -> *). Monad m => [DepGraphAug m] -> DepGraphAug m
applyAugs ([DepGraphAug FusionM] -> DepGraphAug FusionM)
-> [DepGraphAug FusionM] -> DepGraphAug FusionM
forall a b. (a -> b) -> a -> b
$
    [ DepGraphAug FusionM
doSoacThroughTransFusion,
      DepGraphAug FusionM
doVerticalFusion,
      DepGraphAug FusionM
doHorizontalFusion,
      DepGraphAug FusionM
doInnerFusion,
      DepGraphAug FusionM
removeUnusedOutputs
    ]

runInnerFusionOnContext :: DepContext -> FusionM DepContext
runInnerFusionOnContext :: Context NodeT EdgeT -> FusionM (Context NodeT EdgeT)
runInnerFusionOnContext c :: Context NodeT EdgeT
c@(Adj EdgeT
incoming, Int
node, NodeT
nodeT, Adj EdgeT
outgoing) = case NodeT
nodeT of
  DoNode (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Loop [(FParam SOACS, SubExp)]
params LoopForm
form Body SOACS
body)) [(NodeT, [EdgeT])]
to_fuse ->
    FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a. FusionM a -> FusionM a
doFuseScans (FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT))
-> (FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT))
-> FusionM (Context NodeT EdgeT)
-> FusionM (Context NodeT EdgeT)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope SOACS
-> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a. Scope SOACS -> FusionM a -> FusionM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([FParam SOACS] -> Scope SOACS
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (((FParam SOACS, SubExp) -> FParam SOACS)
-> [(FParam SOACS, SubExp)] -> [FParam SOACS]
forall a b. (a -> b) -> [a] -> [b]
map (FParam SOACS, SubExp) -> FParam SOACS
forall a b. (a, b) -> a
fst [(FParam SOACS, SubExp)]
params) Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> LoopForm -> Scope SOACS
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form) (FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT))
-> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a b. (a -> b) -> a -> b
$ do
      b <- Body SOACS -> [(NodeT, [EdgeT])] -> FusionM (Body SOACS)
doFusionWithDelayed Body SOACS
body [(NodeT, [EdgeT])]
to_fuse
      pure (incoming, node, DoNode (Let pat aux (Loop params form b)) [], outgoing)
  MatchNode (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Match [SubExp]
cond [Case (Body SOACS)]
cases Body SOACS
defbody MatchDec (BranchType SOACS)
dec)) [(NodeT, [EdgeT])]
to_fuse -> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a. FusionM a -> FusionM a
doFuseScans (FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT))
-> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a b. (a -> b) -> a -> b
$ do
    cases' <- (Case (Body SOACS) -> FusionM (Case (Body SOACS)))
-> [Case (Body SOACS)] -> FusionM [Case (Body SOACS)]
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 ((Body SOACS -> FusionM (Body SOACS))
-> Case (Body SOACS) -> FusionM (Case (Body SOACS))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Case a -> f (Case b)
traverse ((Body SOACS -> FusionM (Body SOACS))
 -> Case (Body SOACS) -> FusionM (Case (Body SOACS)))
-> (Body SOACS -> FusionM (Body SOACS))
-> Case (Body SOACS)
-> FusionM (Case (Body SOACS))
forall a b. (a -> b) -> a -> b
$ Body SOACS -> FusionM (Body SOACS)
forall rep res (m :: * -> *).
(Renameable rep, Rename res, MonadFreshNames m) =>
GBody rep res -> m (GBody rep res)
renameBody (Body SOACS -> FusionM (Body SOACS))
-> (Body SOACS -> FusionM (Body SOACS))
-> Body SOACS
-> FusionM (Body SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< (Body SOACS -> [(NodeT, [EdgeT])] -> FusionM (Body SOACS)
`doFusionWithDelayed` [(NodeT, [EdgeT])]
to_fuse)) [Case (Body SOACS)]
cases
    defbody' <- doFusionWithDelayed defbody to_fuse
    pure (incoming, node, MatchNode (Let pat aux (Match cond cases' defbody' dec)) [], outgoing)
  StmNode (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op (Futhark.VJP ShapeBase SubExp
shape [SubExp]
args [SubExp]
vec Lambda SOACS
lam))) -> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a. FusionM a -> FusionM a
doFuseScans (FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT))
-> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a b. (a -> b) -> a -> b
$ do
    lam' <- (Lambda SOACS, Bool) -> Lambda SOACS
forall a b. (a, b) -> a
fst ((Lambda SOACS, Bool) -> Lambda SOACS)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
lam
    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.VJP shape args vec lam'))), outgoing)
  StmNode (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op (Futhark.JVP ShapeBase SubExp
shape [SubExp]
args [SubExp]
vec Lambda SOACS
lam))) -> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a. FusionM a -> FusionM a
doFuseScans (FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT))
-> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a b. (a -> b) -> a -> b
$ do
    lam' <- (Lambda SOACS, Bool) -> Lambda SOACS
forall a b. (a, b) -> a
fst ((Lambda SOACS, Bool) -> Lambda SOACS)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
lam
    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.JVP shape args vec lam'))), outgoing)
  StmNode (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
lam)) -> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a. FusionM a -> FusionM a
doFuseScans (FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT))
-> FusionM (Context NodeT EdgeT) -> FusionM (Context NodeT EdgeT)
forall a b. (a -> b) -> a -> b
$ do
    lam' <- (Lambda SOACS, Bool) -> Lambda SOACS
forall a b. (a, b) -> a
fst ((Lambda SOACS, Bool) -> Lambda SOACS)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
lam
    pure (incoming, node, StmNode (Let pat aux (WithAcc inputs lam')), outgoing)
  SoacNode ArrayTransforms
ots Pat Type
pat SOAC SOACS
soac StmAux (ExpDec SOACS)
aux -> do
    soac' <- case SOAC SOACS
soac of
      H.Stream SubExp
w [Input]
inputs [SubExp]
accs Lambda SOACS
lam ->
        SubExp -> [Input] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. SubExp -> [Input] -> [SubExp] -> Lambda rep -> SOAC rep
H.Stream SubExp
w [Input]
inputs [SubExp]
accs (Lambda SOACS -> SOAC SOACS)
-> FusionM (Lambda SOACS) -> FusionM (SOAC SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FusionM (Lambda SOACS) -> FusionM (Lambda SOACS)
forall a. FusionM a -> FusionM a
dontFuseScans (Lambda SOACS -> FusionM (Lambda SOACS)
onLambda Lambda SOACS
lam)
      H.Screma SubExp
w [Input]
inputs (ScremaForm Lambda SOACS
lam [Scan SOACS]
scans [Reduce SOACS]
reds Lambda SOACS
post_lam) ->
        SubExp -> [Input] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [Input] -> ScremaForm rep -> SOAC rep
H.Screma SubExp
w [Input]
inputs
          (ScremaForm SOACS -> SOAC SOACS)
-> FusionM (ScremaForm SOACS) -> FusionM (SOAC SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( Lambda SOACS
-> [Scan SOACS]
-> [Reduce SOACS]
-> Lambda SOACS
-> ScremaForm SOACS
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm
                  (Lambda SOACS
 -> [Scan SOACS]
 -> [Reduce SOACS]
 -> Lambda SOACS
 -> ScremaForm SOACS)
-> FusionM (Lambda SOACS)
-> FusionM
     ([Scan SOACS]
      -> [Reduce SOACS] -> Lambda SOACS -> ScremaForm SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FusionM (Lambda SOACS) -> FusionM (Lambda SOACS)
forall a. FusionM a -> FusionM a
doFuseScans (Lambda SOACS -> FusionM (Lambda SOACS)
onLambda Lambda SOACS
lam)
                  FusionM
  ([Scan SOACS]
   -> [Reduce SOACS] -> Lambda SOACS -> ScremaForm SOACS)
-> FusionM [Scan SOACS]
-> FusionM ([Reduce SOACS] -> Lambda SOACS -> ScremaForm SOACS)
forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Scan SOACS -> FusionM (Scan SOACS))
-> [Scan SOACS] -> FusionM [Scan SOACS]
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 Scan SOACS -> FusionM (Scan SOACS)
onScan [Scan SOACS]
scans
                  FusionM ([Reduce SOACS] -> Lambda SOACS -> ScremaForm SOACS)
-> FusionM [Reduce SOACS]
-> FusionM (Lambda SOACS -> ScremaForm SOACS)
forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Reduce SOACS -> FusionM (Reduce SOACS))
-> [Reduce SOACS] -> FusionM [Reduce SOACS]
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 Reduce SOACS -> FusionM (Reduce SOACS)
onRed [Reduce SOACS]
reds
                  FusionM (Lambda SOACS -> ScremaForm SOACS)
-> FusionM (Lambda SOACS) -> FusionM (ScremaForm SOACS)
forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> FusionM (Lambda SOACS) -> FusionM (Lambda SOACS)
forall a. FusionM a -> FusionM a
doFuseScans (Lambda SOACS -> FusionM (Lambda SOACS)
onLambda Lambda SOACS
post_lam)
              )
      SOAC SOACS
_ ->
        Lambda SOACS -> SOAC SOACS -> SOAC SOACS
forall rep. Lambda rep -> SOAC rep -> SOAC rep
H.setLambda (Lambda SOACS -> SOAC SOACS -> SOAC SOACS)
-> FusionM (Lambda SOACS) -> FusionM (SOAC SOACS -> SOAC SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FusionM (Lambda SOACS) -> FusionM (Lambda SOACS)
forall a. FusionM a -> FusionM a
doFuseScans (Lambda SOACS -> FusionM (Lambda SOACS)
onLambda (SOAC SOACS -> Lambda SOACS
forall rep. SOAC rep -> Lambda rep
H.lambda SOAC SOACS
soac)) FusionM (SOAC SOACS -> SOAC SOACS)
-> FusionM (SOAC SOACS) -> FusionM (SOAC SOACS)
forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOAC SOACS -> FusionM (SOAC SOACS)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SOAC SOACS
soac
    let nodeT' = ArrayTransforms
-> Pat Type -> SOAC SOACS -> StmAux (ExpDec SOACS) -> NodeT
SoacNode ArrayTransforms
ots Pat Type
pat SOAC SOACS
soac' StmAux (ExpDec SOACS)
aux
    pure (incoming, node, nodeT', outgoing)
  NodeT
_ -> Context NodeT EdgeT -> FusionM (Context NodeT EdgeT)
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Context NodeT EdgeT
c
  where
    onLambda :: Lambda SOACS -> FusionM (Lambda SOACS)
onLambda Lambda SOACS
lam = Lambda SOACS -> FusionM (Lambda SOACS) -> FusionM (Lambda SOACS)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Lambda SOACS
lam (FusionM (Lambda SOACS) -> FusionM (Lambda SOACS))
-> (FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS))
-> FusionM (Lambda SOACS, Bool)
-> FusionM (Lambda SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Lambda SOACS, Bool) -> Lambda SOACS)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall a b. (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Lambda SOACS, Bool) -> Lambda SOACS
forall a b. (a, b) -> a
fst (FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS))
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
lam
    onScan :: Scan SOACS -> FusionM (Scan SOACS)
onScan (Scan Lambda SOACS
lam [SubExp]
nes) = Lambda SOACS -> [SubExp] -> Scan SOACS
forall rep. Lambda rep -> [SubExp] -> Scan rep
Scan (Lambda SOACS -> [SubExp] -> Scan SOACS)
-> FusionM (Lambda SOACS) -> FusionM ([SubExp] -> Scan SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS)
onLambda Lambda SOACS
lam FusionM ([SubExp] -> Scan SOACS)
-> FusionM [SubExp] -> FusionM (Scan SOACS)
forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [SubExp] -> FusionM [SubExp]
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
nes
    onRed :: Reduce SOACS -> FusionM (Reduce SOACS)
onRed (Reduce Commutativity
comm Lambda SOACS
lam [SubExp]
nes) = Commutativity -> Lambda SOACS -> [SubExp] -> Reduce SOACS
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
comm (Lambda SOACS -> [SubExp] -> Reduce SOACS)
-> FusionM (Lambda SOACS) -> FusionM ([SubExp] -> Reduce SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> FusionM (Lambda SOACS)
onLambda Lambda SOACS
lam FusionM ([SubExp] -> Reduce SOACS)
-> FusionM [SubExp] -> FusionM (Reduce SOACS)
forall a b. FusionM (a -> b) -> FusionM a -> FusionM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [SubExp] -> FusionM [SubExp]
forall a. a -> FusionM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
nes

    doFusionWithDelayed :: Body SOACS -> [(NodeT, [EdgeT])] -> FusionM (Body SOACS)
    doFusionWithDelayed :: Body SOACS -> [(NodeT, [EdgeT])] -> FusionM (Body SOACS)
doFusionWithDelayed (Body () Stms SOACS
stms [SubExpRes]
res) [(NodeT, [EdgeT])]
extraNodes = Stms SOACS -> FusionM (Body SOACS) -> FusionM (Body SOACS)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms SOACS
stms (FusionM (Body SOACS) -> FusionM (Body SOACS))
-> FusionM (Body SOACS) -> FusionM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
      stm_node <- ((NodeT, [EdgeT]) -> FusionM (Stms SOACS))
-> [(NodeT, [EdgeT])] -> FusionM [Stms SOACS]
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 (NodeT -> FusionM (Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
NodeT -> m (Stms SOACS)
finalizeNode (NodeT -> FusionM (Stms SOACS))
-> ((NodeT, [EdgeT]) -> NodeT)
-> (NodeT, [EdgeT])
-> FusionM (Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NodeT, [EdgeT]) -> NodeT
forall a b. (a, b) -> a
fst) [(NodeT, [EdgeT])]
extraNodes
      stms' <- fuseGraph (mkBody (mconcat stm_node <> stms) res)
      pure $ Body () stms' res

doFusionInLambda :: Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda :: Lambda SOACS -> FusionM (Lambda SOACS, Bool)
doFusionInLambda Lambda SOACS
lam = do
  (Lambda SOACS, Bool)
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS, Bool)
forall a. a -> FusionM a -> FusionM a
useGas (Lambda SOACS
lam, Bool
False) (FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS, Bool))
-> FusionM (Lambda SOACS, Bool) -> FusionM (Lambda SOACS, Bool)
forall a b. (a -> b) -> a -> b
$ do
    -- To clean up previous instances of fusion.
    lam' <- Lambda SOACS -> FusionM (Lambda SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Lambda SOACS -> m (Lambda SOACS)
simplifyLambda Lambda SOACS
lam
    prev_count <- gets fusionCount
    newbody <- inScopeOf lam' $ doFusionBody $ lambdaBody lam'
    aft_count <- gets fusionCount
    -- To clean up any inner fusion.
    lam'' <-
      (if prev_count /= aft_count then simplifyLambda else pure)
        lam' {lambdaBody = newbody}
    pure (lam'', prev_count /= aft_count)
  where
    doFusionBody :: Body SOACS -> FusionM (Body SOACS)
    doFusionBody :: Body SOACS -> FusionM (Body SOACS)
doFusionBody Body SOACS
body = do
      stms' <- Body SOACS -> FusionM (Stms SOACS)
fuseGraph Body SOACS
body
      pure $ body {bodyStms = stms'}

-- main fusion function.
fuseGraph :: Body SOACS -> FusionM (Stms SOACS)
fuseGraph :: Body SOACS -> FusionM (Stms SOACS)
fuseGraph Body SOACS
body = Stms SOACS -> FusionM (Stms SOACS) -> FusionM (Stms SOACS)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf (Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body) (FusionM (Stms SOACS) -> FusionM (Stms SOACS))
-> FusionM (Stms SOACS) -> FusionM (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ do
  graph_not_fused <- Body SOACS -> FusionM DepGraph
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Body SOACS -> m DepGraph
mkDepGraph Body SOACS
body
  graph_fused <- doAllFusion graph_not_fused
  linearizeGraph graph_fused

fuseConsts :: Maybe Int -> [VName] -> Stms SOACS -> PassM (Stms SOACS)
fuseConsts :: Maybe Int -> [VName] -> Stms SOACS -> PassM (Stms SOACS)
fuseConsts Maybe Int
g [VName]
outputs Stms SOACS
stms =
  Scope SOACS
-> FusionEnv -> FusionM (Stms SOACS) -> PassM (Stms SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> FusionEnv -> FusionM a -> m a
runFusionM
    (Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
stms)
    (Maybe Int -> FusionEnv
freshFusionEnv Maybe Int
g)
    (Body SOACS -> FusionM (Stms SOACS)
fuseGraph (Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
stms ([VName] -> [SubExpRes]
varsRes [VName]
outputs)))

fuseFun :: Maybe Int -> Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)
fuseFun :: Maybe Int -> Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)
fuseFun Maybe Int
g Stms SOACS
consts FunDef SOACS
fun = do
  fun_stms' <-
    Scope SOACS
-> FusionEnv -> FusionM (Stms SOACS) -> PassM (Stms SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> FusionEnv -> FusionM a -> m a
runFusionM
      (FunDef SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf FunDef SOACS
fun Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
consts)
      (Maybe Int -> FusionEnv
freshFusionEnv Maybe Int
g)
      (Body SOACS -> FusionM (Stms SOACS)
fuseGraph (FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fun))
  pure fun {funDefBody = (funDefBody fun) {bodyStms = fun_stms'}}

-- | Pass definition with an optional bound on the number of
-- iterations of the fusion convergence loop.  'Just n' runs at
-- most @n@ iterations; 'Nothing' means no bound (iterate until convergence).
{-# NOINLINE fuseSOACs #-}
fuseSOACs :: Maybe Int -> Pass SOACS SOACS
fuseSOACs :: Maybe Int -> Pass SOACS SOACS
fuseSOACs Maybe Int
g =
  Pass
    { passName :: String
passName = String
"Fuse SOACs",
      passDescription :: String
passDescription = String
"Perform higher-order optimisation, i.e., fusion.",
      passFunction :: Prog SOACS -> PassM (Prog SOACS)
passFunction = \Prog SOACS
p ->
        (Stms SOACS -> PassM (Stms SOACS))
-> (Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS))
-> Prog SOACS
-> PassM (Prog SOACS)
forall fromrep torep.
(Stms fromrep -> PassM (Stms torep))
-> (Stms torep -> FunDef fromrep -> PassM (FunDef torep))
-> Prog fromrep
-> PassM (Prog torep)
intraproceduralTransformationWithConsts
          (Maybe Int -> [VName] -> Stms SOACS -> PassM (Stms SOACS)
fuseConsts Maybe Int
g (Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ [FunDef SOACS] -> Names
forall a. FreeIn a => a -> Names
freeIn (Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
p)))
          (Maybe Int -> Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)
fuseFun Maybe Int
g)
          Prog SOACS
p
    }