-- | A graph representation of a sequence of Futhark statements
-- (i.e. a 'Body'), built to handle fusion.  Could perhaps be made
-- more general.  An important property is that it does not handle
-- "nested bodies" (e.g. 'Match'); these are represented as single
-- nodes.
--
-- This is all implemented on top of the graph representation provided
-- by the @fgl@ package ("Data.Graph.Inductive").  The graph provided
-- by this package allows nodes and edges to have arbitrarily-typed
-- "labels".  It is these labels ('EdgeT', 'NodeT') that we use to
-- contain Futhark-specific information.  An edge goes *from*
-- consumers to producers.  There are also edges that do not represent
-- normal data dependencies, but other things.  This means that a node
-- can have multiple edges for the same name, indicating different
-- kinds of dependencies.
module Futhark.Optimise.Fusion.GraphRep
  ( -- * Data structure
    EdgeT (..),
    NodeT (..),
    DepContext,
    DepGraphAug,
    DepGraph (..),
    DepNode,

    -- * Queries
    getName,
    nodeFromLNode,
    mergedContext,
    mapAcross,
    edgesBetween,
    reachable,
    applyAugs,
    depsFromEdge,
    contractEdge,
    isRealNode,
    isCons,
    isDep,
    isInf,

    -- * Construction
    mkDepGraph,
    mkDepGraphForFun,
    pprg,
    isWithAccNodeT,
    isWithAccNodeId,
    vFusionFeasability,
    hFusionFeasability,
  )
where

import Control.Monad.Reader
import Data.Bifunctor (bimap)
import Data.Foldable (foldlM)
import Data.Graph.Inductive.Dot qualified as G
import Data.Graph.Inductive.Graph qualified as G
import Data.Graph.Inductive.Query.DFS qualified as Q
import Data.Graph.Inductive.Tree qualified as G
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe (mapMaybe)
import Data.Set qualified as S
import Futhark.Analysis.Alias qualified as Alias
import Futhark.Analysis.HORep.SOAC qualified as H
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS hiding (SOAC (..))
import Futhark.IR.SOACS qualified as Futhark
import Futhark.Util (nubOrd)

-- | Information associated with an edge in the graph.
data EdgeT
  = Alias VName
  | InfDep VName
  | Dep VName
  | Cons VName
  | Fake VName
  | Res VName
  deriving (EdgeT -> EdgeT -> Bool
(EdgeT -> EdgeT -> Bool) -> (EdgeT -> EdgeT -> Bool) -> Eq EdgeT
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EdgeT -> EdgeT -> Bool
== :: EdgeT -> EdgeT -> Bool
$c/= :: EdgeT -> EdgeT -> Bool
/= :: EdgeT -> EdgeT -> Bool
Eq, Eq EdgeT
Eq EdgeT =>
(EdgeT -> EdgeT -> Ordering)
-> (EdgeT -> EdgeT -> Bool)
-> (EdgeT -> EdgeT -> Bool)
-> (EdgeT -> EdgeT -> Bool)
-> (EdgeT -> EdgeT -> Bool)
-> (EdgeT -> EdgeT -> EdgeT)
-> (EdgeT -> EdgeT -> EdgeT)
-> Ord EdgeT
EdgeT -> EdgeT -> Bool
EdgeT -> EdgeT -> Ordering
EdgeT -> EdgeT -> EdgeT
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: EdgeT -> EdgeT -> Ordering
compare :: EdgeT -> EdgeT -> Ordering
$c< :: EdgeT -> EdgeT -> Bool
< :: EdgeT -> EdgeT -> Bool
$c<= :: EdgeT -> EdgeT -> Bool
<= :: EdgeT -> EdgeT -> Bool
$c> :: EdgeT -> EdgeT -> Bool
> :: EdgeT -> EdgeT -> Bool
$c>= :: EdgeT -> EdgeT -> Bool
>= :: EdgeT -> EdgeT -> Bool
$cmax :: EdgeT -> EdgeT -> EdgeT
max :: EdgeT -> EdgeT -> EdgeT
$cmin :: EdgeT -> EdgeT -> EdgeT
min :: EdgeT -> EdgeT -> EdgeT
Ord)

-- | Information associated with a node in the graph.
data NodeT
  = StmNode (Stm SOACS)
  | SoacNode H.ArrayTransforms (Pat Type) (H.SOAC SOACS) (StmAux (ExpDec SOACS))
  | -- | First 'VName' is result; last is input.
    TransNode VName H.ArrayTransform VName
  | -- | Node corresponding to a result of the entire computation
    -- (i.e. the 'Result' of a body).  Any node that is not
    -- transitively reachable from one of these can be considered
    -- dead.
    ResNode VName
  | -- | Node corresponding to a free variable.  These are used to
    -- safely handle consumption, which also means we don't have to
    -- create a node for every free single variable.
    FreeNode VName
  | MatchNode (Stm SOACS) [(NodeT, [EdgeT])]
  | DoNode (Stm SOACS) [(NodeT, [EdgeT])]
  deriving (NodeT -> NodeT -> Bool
(NodeT -> NodeT -> Bool) -> (NodeT -> NodeT -> Bool) -> Eq NodeT
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NodeT -> NodeT -> Bool
== :: NodeT -> NodeT -> Bool
$c/= :: NodeT -> NodeT -> Bool
/= :: NodeT -> NodeT -> Bool
Eq)

instance Show EdgeT where
  show :: EdgeT -> String
show (Dep VName
vName) = String
"Dep " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> VName -> String
forall a. Pretty a => a -> String
prettyString VName
vName
  show (InfDep VName
vName) = String
"iDep " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> VName -> String
forall a. Pretty a => a -> String
prettyString VName
vName
  show (Cons VName
_) = String
"Cons"
  show (Fake VName
_) = String
"Fake"
  show (Res VName
_) = String
"Res"
  show (Alias VName
_) = String
"Alias"

instance Show NodeT where
  show :: NodeT -> String
show (StmNode (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
_ Exp SOACS
_)) = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
L.intercalate String
", " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (VName -> String) -> [VName] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map VName -> String
forall a. Pretty a => a -> String
prettyString ([VName] -> [String]) -> [VName] -> [String]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
Pat (LetDec SOACS)
pat
  show (SoacNode ArrayTransforms
_ Pat Type
pat SOAC SOACS
_ StmAux (ExpDec SOACS)
_) = Pat Type -> String
forall a. Pretty a => a -> String
prettyString Pat Type
pat
  show (TransNode VName
_ ArrayTransform
tr VName
_) = ArrayTransform -> String
forall a. Pretty a => a -> String
prettyString ArrayTransform
tr
  show (ResNode VName
name) = ShowS
forall a. Pretty a => a -> String
prettyString ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String
"Res: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
name
  show (FreeNode VName
name) = ShowS
forall a. Pretty a => a -> String
prettyString ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String
"Input: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
name
  show (MatchNode Stm SOACS
stm [(NodeT, [EdgeT])]
_) = String
"Match: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
L.intercalate String
", " ((VName -> String) -> [VName] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map VName -> String
forall a. Pretty a => a -> String
prettyString ([VName] -> [String]) -> [VName] -> [String]
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> [VName]
stmNames Stm SOACS
stm)
  show (DoNode Stm SOACS
stm [(NodeT, [EdgeT])]
_) = String
"Do: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
L.intercalate String
", " ((VName -> String) -> [VName] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map VName -> String
forall a. Pretty a => a -> String
prettyString ([VName] -> [String]) -> [VName] -> [String]
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> [VName]
stmNames Stm SOACS
stm)

-- | The name that this edge depends on.
getName :: EdgeT -> VName
getName :: EdgeT -> VName
getName EdgeT
edgeT = case EdgeT
edgeT of
  Alias VName
vn -> VName
vn
  InfDep VName
vn -> VName
vn
  Dep VName
vn -> VName
vn
  Cons VName
vn -> VName
vn
  Fake VName
vn -> VName
vn
  Res VName
vn -> VName
vn

-- | Does the node acutally represent something in the program?  A
-- "non-real" node represents things like fake nodes inserted to
-- express ordering due to consumption.
isRealNode :: NodeT -> Bool
isRealNode :: NodeT -> Bool
isRealNode ResNode {} = Bool
False
isRealNode FreeNode {} = Bool
False
isRealNode NodeT
_ = Bool
True

-- | Prettyprint dependency graph.
pprg :: DepGraph -> String
pprg :: DepGraph -> String
pprg = Dot () -> String
forall a. Dot a -> String
G.showDot (Dot () -> String) -> (DepGraph -> Dot ()) -> DepGraph -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gr String String -> Dot ()
forall (gr :: * -> * -> *). Graph gr => gr String String -> Dot ()
G.fglToDotString (Gr String String -> Dot ())
-> (DepGraph -> Gr String String) -> DepGraph -> Dot ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NodeT -> String)
-> (EdgeT -> String) -> Gr NodeT EdgeT -> Gr String String
forall (gr :: * -> * -> *) a c b d.
DynGraph gr =>
(a -> c) -> (b -> d) -> gr a b -> gr c d
G.nemap NodeT -> String
forall a. Show a => a -> String
show EdgeT -> String
forall a. Show a => a -> String
show (Gr NodeT EdgeT -> Gr String String)
-> (DepGraph -> Gr NodeT EdgeT) -> DepGraph -> Gr String String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DepGraph -> Gr NodeT EdgeT
dgGraph

-- | A pair of a 'G.Node' and the node label.
type DepNode = G.LNode NodeT

type DepEdge = G.LEdge EdgeT

-- | A tuple with four parts: inbound links to the node, the node
-- itself, the 'NodeT' "label", and outbound links from the node.
-- This type is used to modify the graph in 'mapAcross'.
type DepContext = G.Context NodeT EdgeT

-- | A dependency graph.  Edges go from *consumers* to *producers*
-- (i.e. from usage to definition).  That means the incoming edges of
-- a node are the dependents of that node, and the outgoing edges are
-- the dependencies of that node.
data DepGraph = DepGraph
  { DepGraph -> Gr NodeT EdgeT
dgGraph :: G.Gr NodeT EdgeT,
    DepGraph -> ProducerMapping
dgProducerMapping :: ProducerMapping,
    -- | A table mapping VNames to VNames that are aliased to it.
    DepGraph -> AliasTable
dgAliasTable :: AliasTable
  }

-- | A "graph augmentation" is a monadic action that modifies the graph.
type DepGraphAug m = DepGraph -> m DepGraph

-- | For each node, what producer should the node depend on and what
-- type is it.
type EdgeGenerator = NodeT -> [(VName, EdgeT)]

-- | A mapping from variable name to the graph node that produces
-- it.
type ProducerMapping = M.Map VName G.Node

makeMapping :: (Monad m) => DepGraphAug m
makeMapping :: forall (m :: * -> *). Monad m => DepGraphAug m
makeMapping dg :: DepGraph
dg@(DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
dgGraph = Gr NodeT EdgeT
g}) =
  DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg {dgProducerMapping = M.fromList $ concatMap gen_dep_list (G.labNodes g)}
  where
    gen_dep_list :: DepNode -> [(VName, G.Node)]
    gen_dep_list :: DepNode -> [(VName, Node)]
gen_dep_list (Node
i, NodeT
node) = [(VName
name, Node
i) | VName
name <- NodeT -> [VName]
getOutputs NodeT
node]

-- | Apply several graph augmentations in sequence.
applyAugs :: (Monad m) => [DepGraphAug m] -> DepGraphAug m
applyAugs :: forall (m :: * -> *). Monad m => [DepGraphAug m] -> DepGraphAug m
applyAugs [DepGraphAug m]
augs DepGraph
g = (DepGraph -> DepGraphAug m -> m DepGraph)
-> DepGraph -> [DepGraphAug m] -> m DepGraph
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM ((DepGraphAug m -> DepGraphAug m)
-> DepGraph -> DepGraphAug m -> m DepGraph
forall a b c. (a -> b -> c) -> b -> a -> c
flip DepGraphAug m -> DepGraphAug m
forall a b. (a -> b) -> a -> b
($)) DepGraph
g [DepGraphAug m]
augs

-- | Creates deps for the given nodes on the graph using the 'EdgeGenerator'.
genEdges :: (Monad m) => [DepNode] -> EdgeGenerator -> DepGraphAug m
genEdges :: forall (m :: * -> *).
Monad m =>
[DepNode] -> EdgeGenerator -> DepGraphAug m
genEdges [DepNode]
l_stms EdgeGenerator
edge_fun DepGraph
dg =
  [DepEdge] -> DepGraphAug m
forall (m :: * -> *). Monad m => [DepEdge] -> DepGraphAug m
depGraphInsertEdges ((DepNode -> [DepEdge]) -> [DepNode] -> [DepEdge]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (ProducerMapping -> DepNode -> [DepEdge]
genEdge (DepGraph -> ProducerMapping
dgProducerMapping DepGraph
dg)) [DepNode]
l_stms) DepGraph
dg
  where
    -- statements -> mapping from declared array names to soac index
    genEdge :: M.Map VName G.Node -> DepNode -> [G.LEdge EdgeT]
    genEdge :: ProducerMapping -> DepNode -> [DepEdge]
genEdge ProducerMapping
name_map (Node
from, NodeT
node) = do
      (VName
dep, EdgeT
edgeT) <- EdgeGenerator
edge_fun NodeT
node
      Just Node
to <- [VName -> ProducerMapping -> Maybe Node
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
dep ProducerMapping
name_map]
      DepEdge -> [DepEdge]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DepEdge -> [DepEdge]) -> DepEdge -> [DepEdge]
forall a b. (a -> b) -> a -> b
$ Edge -> EdgeT -> DepEdge
forall b. Edge -> b -> LEdge b
G.toLEdge (Node
from, Node
to) EdgeT
edgeT

depGraphInsertEdges :: (Monad m) => [DepEdge] -> DepGraphAug m
depGraphInsertEdges :: forall (m :: * -> *). Monad m => [DepEdge] -> DepGraphAug m
depGraphInsertEdges [DepEdge]
edgs DepGraph
dg = DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DepGraph -> m DepGraph) -> DepGraph -> m DepGraph
forall a b. (a -> b) -> a -> b
$ DepGraph
dg {dgGraph = G.insEdges edgs $ dgGraph dg}

-- | Monadically modify every node of the graph.
mapAcross :: (Monad m) => (DepContext -> m DepContext) -> DepGraphAug m
mapAcross :: forall (m :: * -> *).
Monad m =>
(DepContext -> m DepContext) -> DepGraphAug m
mapAcross DepContext -> m DepContext
f DepGraph
dg = do
  Gr NodeT EdgeT
g' <- (Gr NodeT EdgeT -> Node -> m (Gr NodeT EdgeT))
-> Gr NodeT EdgeT -> [Node] -> m (Gr NodeT EdgeT)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM ((Node -> Gr NodeT EdgeT -> m (Gr NodeT EdgeT))
-> Gr NodeT EdgeT -> Node -> m (Gr NodeT EdgeT)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Node -> Gr NodeT EdgeT -> m (Gr NodeT EdgeT)
forall {gr :: * -> * -> *}.
DynGraph gr =>
Node -> gr NodeT EdgeT -> m (gr NodeT EdgeT)
helper) (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg) (Gr NodeT EdgeT -> [Node]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [Node]
G.nodes (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg))
  DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DepGraph -> m DepGraph) -> DepGraph -> m DepGraph
forall a b. (a -> b) -> a -> b
$ DepGraph
dg {dgGraph = g'}
  where
    helper :: Node -> gr NodeT EdgeT -> m (gr NodeT EdgeT)
helper Node
n gr NodeT EdgeT
g' = case Node -> gr NodeT EdgeT -> Decomp gr NodeT EdgeT
forall a b. Node -> gr a b -> Decomp gr a b
forall (gr :: * -> * -> *) a b.
Graph gr =>
Node -> gr a b -> Decomp gr a b
G.match Node
n gr NodeT EdgeT
g' of
      (Just DepContext
c, gr NodeT EdgeT
g_new) -> do
        DepContext
c' <- DepContext -> m DepContext
f DepContext
c
        gr NodeT EdgeT -> m (gr NodeT EdgeT)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (gr NodeT EdgeT -> m (gr NodeT EdgeT))
-> gr NodeT EdgeT -> m (gr NodeT EdgeT)
forall a b. (a -> b) -> a -> b
$ DepContext
c' DepContext -> gr NodeT EdgeT -> gr NodeT EdgeT
forall a b. Context a b -> gr a b -> gr a b
forall (gr :: * -> * -> *) a b.
DynGraph gr =>
Context a b -> gr a b -> gr a b
G.& gr NodeT EdgeT
g_new
      (Maybe DepContext
Nothing, gr NodeT EdgeT
_) -> gr NodeT EdgeT -> m (gr NodeT EdgeT)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure gr NodeT EdgeT
g'

stmFromNode :: NodeT -> Stms SOACS -- do not use outside of edge generation
stmFromNode :: NodeT -> Stms SOACS
stmFromNode (StmNode Stm SOACS
x) = Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
x
stmFromNode NodeT
_ = Stms SOACS
forall a. Monoid a => a
mempty

-- | Get the underlying @fgl@ node.
nodeFromLNode :: DepNode -> G.Node
nodeFromLNode :: DepNode -> Node
nodeFromLNode = DepNode -> Node
forall a b. (a, b) -> a
fst

-- | Get the variable name that this edge refers to.
depsFromEdge :: DepEdge -> VName
depsFromEdge :: DepEdge -> VName
depsFromEdge = EdgeT -> VName
getName (EdgeT -> VName) -> (DepEdge -> EdgeT) -> DepEdge -> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DepEdge -> EdgeT
forall b. LEdge b -> b
G.edgeLabel

-- | Find all the edges connecting the two nodes.
edgesBetween :: DepGraph -> G.Node -> G.Node -> [DepEdge]
edgesBetween :: DepGraph -> Node -> Node -> [DepEdge]
edgesBetween DepGraph
dg Node
n1 Node
n2 = Gr NodeT EdgeT -> [DepEdge]
forall a b. Gr a b -> [LEdge b]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LEdge b]
G.labEdges (Gr NodeT EdgeT -> [DepEdge]) -> Gr NodeT EdgeT -> [DepEdge]
forall a b. (a -> b) -> a -> b
$ [Node] -> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall (gr :: * -> * -> *) a b.
DynGraph gr =>
[Node] -> gr a b -> gr a b
G.subgraph [Node
n1, Node
n2] (Gr NodeT EdgeT -> Gr NodeT EdgeT)
-> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall a b. (a -> b) -> a -> b
$ DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg

-- | @reachable dg from to@ is true if @to@ is reachable from @from@.
reachable :: DepGraph -> G.Node -> G.Node -> Bool
reachable :: DepGraph -> Node -> Node -> Bool
reachable DepGraph
dg Node
source Node
target = Node
target Node -> [Node] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Node -> Gr NodeT EdgeT -> [Node]
forall (gr :: * -> * -> *) a b.
Graph gr =>
Node -> gr a b -> [Node]
Q.reachable Node
source (DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg)

-- Utility func for augs
augWithFun :: (Monad m) => EdgeGenerator -> DepGraphAug m
augWithFun :: forall (m :: * -> *). Monad m => EdgeGenerator -> DepGraphAug m
augWithFun EdgeGenerator
f DepGraph
dg = [DepNode] -> EdgeGenerator -> DepGraphAug m
forall (m :: * -> *).
Monad m =>
[DepNode] -> EdgeGenerator -> DepGraphAug m
genEdges (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)) EdgeGenerator
f DepGraph
dg

addDeps :: (Monad m) => DepGraphAug m
addDeps :: forall (m :: * -> *). Monad m => DepGraphAug m
addDeps = EdgeGenerator -> DepGraphAug m
forall (m :: * -> *). Monad m => EdgeGenerator -> DepGraphAug m
augWithFun EdgeGenerator
toDep
  where
    toDep :: EdgeGenerator
toDep NodeT
stmt =
      let ([VName]
fusible, [VName]
infusible) =
            ([(VName, Classification)] -> [VName])
-> ([(VName, Classification)] -> [VName])
-> ([(VName, Classification)], [(VName, Classification)])
-> ([VName], [VName])
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (((VName, Classification) -> VName)
-> [(VName, Classification)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Classification) -> VName
forall a b. (a, b) -> a
fst) (((VName, Classification) -> VName)
-> [(VName, Classification)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Classification) -> VName
forall a b. (a, b) -> a
fst)
              (([(VName, Classification)], [(VName, Classification)])
 -> ([VName], [VName]))
-> (Classifications
    -> ([(VName, Classification)], [(VName, Classification)]))
-> Classifications
-> ([VName], [VName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, Classification) -> Bool)
-> [(VName, Classification)]
-> ([(VName, Classification)], [(VName, Classification)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition ((Classification -> Classification -> Bool
forall a. Eq a => a -> a -> Bool
== Classification
SOACInput) (Classification -> Bool)
-> ((VName, Classification) -> Classification)
-> (VName, Classification)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Classification) -> Classification
forall a b. (a, b) -> b
snd)
              ([(VName, Classification)]
 -> ([(VName, Classification)], [(VName, Classification)]))
-> (Classifications -> [(VName, Classification)])
-> Classifications
-> ([(VName, Classification)], [(VName, Classification)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Classifications -> [(VName, Classification)]
forall a. Set a -> [a]
S.toList
              (Classifications -> ([VName], [VName]))
-> Classifications -> ([VName], [VName])
forall a b. (a -> b) -> a -> b
$ (Stm SOACS -> Classifications) -> Stms SOACS -> Classifications
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm SOACS -> Classifications
stmInputs (NodeT -> Stms SOACS
stmFromNode NodeT
stmt)
          mkDep :: VName -> (VName, EdgeT)
mkDep VName
vname = (VName
vname, VName -> EdgeT
Dep VName
vname)
          mkInfDep :: VName -> (VName, EdgeT)
mkInfDep VName
vname = (VName
vname, VName -> EdgeT
InfDep VName
vname)
       in (VName -> (VName, EdgeT)) -> [VName] -> [(VName, EdgeT)]
forall a b. (a -> b) -> [a] -> [b]
map VName -> (VName, EdgeT)
mkDep [VName]
fusible [(VName, EdgeT)] -> [(VName, EdgeT)] -> [(VName, EdgeT)]
forall a. Semigroup a => a -> a -> a
<> (VName -> (VName, EdgeT)) -> [VName] -> [(VName, EdgeT)]
forall a b. (a -> b) -> [a] -> [b]
map VName -> (VName, EdgeT)
mkInfDep [VName]
infusible

addConsAndAliases :: (Monad m) => DepGraphAug m
addConsAndAliases :: forall (m :: * -> *). Monad m => DepGraphAug m
addConsAndAliases = EdgeGenerator -> DepGraphAug m
forall (m :: * -> *). Monad m => EdgeGenerator -> DepGraphAug m
augWithFun EdgeGenerator
edges
  where
    edges :: EdgeGenerator
edges (StmNode Stm SOACS
s) = Stm (Aliases SOACS) -> [(VName, EdgeT)]
forall {rep}. Aliased rep => Stm rep -> [(VName, EdgeT)]
consEdges Stm (Aliases SOACS)
s' [(VName, EdgeT)] -> [(VName, EdgeT)] -> [(VName, EdgeT)]
forall a. Semigroup a => a -> a -> a
<> Stm (Aliases SOACS) -> [(VName, EdgeT)]
aliasEdges Stm (Aliases SOACS)
s'
      where
        s' :: Stm (Aliases SOACS)
s' = AliasTable -> Stm SOACS -> Stm (Aliases SOACS)
forall rep.
AliasableRep rep =>
AliasTable -> Stm rep -> Stm (Aliases rep)
Alias.analyseStm AliasTable
forall a. Monoid a => a
mempty Stm SOACS
s
    edges NodeT
_ = [(VName, EdgeT)]
forall a. Monoid a => a
mempty
    consEdges :: Stm rep -> [(VName, EdgeT)]
consEdges Stm rep
s = [VName] -> [EdgeT] -> [(VName, EdgeT)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names ((VName -> EdgeT) -> [VName] -> [EdgeT]
forall a b. (a -> b) -> [a] -> [b]
map VName -> EdgeT
Cons [VName]
names)
      where
        names :: [VName]
names = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm rep -> Names
forall rep. Aliased rep => Stm rep -> Names
consumedInStm Stm rep
s
    aliasEdges :: Stm (Aliases SOACS) -> [(VName, EdgeT)]
aliasEdges =
      (VName -> (VName, EdgeT)) -> [VName] -> [(VName, EdgeT)]
forall a b. (a -> b) -> [a] -> [b]
map (\VName
vname -> (VName
vname, VName -> EdgeT
Alias VName
vname))
        ([VName] -> [(VName, EdgeT)])
-> (Stm (Aliases SOACS) -> [VName])
-> Stm (Aliases SOACS)
-> [(VName, EdgeT)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList
        (Names -> [VName])
-> (Stm (Aliases SOACS) -> Names) -> Stm (Aliases SOACS) -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat
        ([Names] -> Names)
-> (Stm (Aliases SOACS) -> [Names]) -> Stm (Aliases SOACS) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec (Aliases SOACS)) -> [Names]
forall dec. AliasesOf dec => Pat dec -> [Names]
patAliases
        (Pat (LetDec (Aliases SOACS)) -> [Names])
-> (Stm (Aliases SOACS) -> Pat (LetDec (Aliases SOACS)))
-> Stm (Aliases SOACS)
-> [Names]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm (Aliases SOACS) -> Pat (LetDec (Aliases SOACS))
forall rep. Stm rep -> Pat (LetDec rep)
stmPat

-- extra dependencies mask the fact that consuming nodes "depend" on all other
-- nodes coming before it (now also adds fake edges to aliases - hope this
-- fixes asymptotic complexity guarantees)
addExtraCons :: (Monad m) => DepGraphAug m
addExtraCons :: forall (m :: * -> *). Monad m => DepGraphAug m
addExtraCons DepGraph
dg =
  [DepEdge] -> DepGraphAug m
forall (m :: * -> *). Monad m => [DepEdge] -> DepGraphAug m
depGraphInsertEdges ((DepEdge -> [DepEdge]) -> [DepEdge] -> [DepEdge]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DepEdge -> [DepEdge]
makeEdge (Gr NodeT EdgeT -> [DepEdge]
forall a b. Gr a b -> [LEdge b]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LEdge b]
G.labEdges Gr NodeT EdgeT
g)) DepGraph
dg
  where
    g :: Gr NodeT EdgeT
g = DepGraph -> Gr NodeT EdgeT
dgGraph DepGraph
dg
    alias_table :: AliasTable
alias_table = DepGraph -> AliasTable
dgAliasTable DepGraph
dg
    mapping :: ProducerMapping
mapping = DepGraph -> ProducerMapping
dgProducerMapping DepGraph
dg
    makeEdge :: DepEdge -> [DepEdge]
makeEdge (Node
from, Node
to, Cons VName
cname) = do
      let aliases :: [VName]
aliases = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Names -> VName -> AliasTable -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
cname AliasTable
alias_table
          to' :: [Node]
to' = (VName -> Maybe Node) -> [VName] -> [Node]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName -> ProducerMapping -> Maybe Node
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` ProducerMapping
mapping) [VName]
aliases
          p :: (Node, EdgeT) -> Bool
p (Node
tonode, EdgeT
toedge) =
            Node
tonode Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
/= Node
from Bool -> Bool -> Bool
&& EdgeT -> VName
getName EdgeT
toedge VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` (VName
cname VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
aliases)
      (Node
to2, EdgeT
_) <- ((Node, EdgeT) -> Bool) -> [(Node, EdgeT)] -> [(Node, EdgeT)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Node, EdgeT) -> Bool
p ([(Node, EdgeT)] -> [(Node, EdgeT)])
-> [(Node, EdgeT)] -> [(Node, EdgeT)]
forall a b. (a -> b) -> a -> b
$ (Node -> [(Node, EdgeT)]) -> [Node] -> [(Node, EdgeT)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Gr NodeT EdgeT -> Node -> [(Node, EdgeT)]
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Node -> [(Node, b)]
G.lpre Gr NodeT EdgeT
g) [Node]
to' [(Node, EdgeT)] -> [(Node, EdgeT)] -> [(Node, EdgeT)]
forall a. Semigroup a => a -> a -> a
<> Gr NodeT EdgeT -> Node -> [(Node, EdgeT)]
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Node -> [(Node, b)]
G.lpre Gr NodeT EdgeT
g Node
to
      DepEdge -> [DepEdge]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DepEdge -> [DepEdge]) -> DepEdge -> [DepEdge]
forall a b. (a -> b) -> a -> b
$ Edge -> EdgeT -> DepEdge
forall b. Edge -> b -> LEdge b
G.toLEdge (Node
from, Node
to2) (VName -> EdgeT
Fake VName
cname)
    makeEdge DepEdge
_ = []

mapAcrossNodeTs :: (Monad m) => (NodeT -> m NodeT) -> DepGraphAug m
mapAcrossNodeTs :: forall (m :: * -> *).
Monad m =>
(NodeT -> m NodeT) -> DepGraphAug m
mapAcrossNodeTs NodeT -> m NodeT
f = (DepContext -> m DepContext) -> DepGraphAug m
forall (m :: * -> *).
Monad m =>
(DepContext -> m DepContext) -> DepGraphAug m
mapAcross DepContext -> m DepContext
forall {a} {b} {d}. (a, b, NodeT, d) -> m (a, b, NodeT, d)
f'
  where
    f' :: (a, b, NodeT, d) -> m (a, b, NodeT, d)
f' (a
ins, b
n, NodeT
nodeT, d
outs) = do
      NodeT
nodeT' <- NodeT -> m NodeT
f NodeT
nodeT
      (a, b, NodeT, d) -> m (a, b, NodeT, d)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
ins, b
n, NodeT
nodeT', d
outs)

nodeToSoacNode :: (HasScope SOACS m, Monad m) => NodeT -> m NodeT
nodeToSoacNode :: forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
NodeT -> m NodeT
nodeToSoacNode n :: NodeT
n@(StmNode s :: Stm SOACS
s@(Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
op)) = case Exp SOACS
op of
  Op {} -> do
    Either NotSOAC (SOAC SOACS)
maybeSoac <- Exp SOACS -> m (Either NotSOAC (SOAC SOACS))
forall rep (m :: * -> *).
(Op rep ~ SOAC rep, HasScope rep m) =>
Exp rep -> m (Either NotSOAC (SOAC rep))
H.fromExp Exp SOACS
op
    case Either NotSOAC (SOAC SOACS)
maybeSoac of
      Right SOAC SOACS
hsoac -> NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT -> m NodeT) -> NodeT -> m 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 Pat Type
Pat (LetDec SOACS)
pat SOAC SOACS
hsoac StmAux (ExpDec SOACS)
aux
      Left NotSOAC
H.NotSOAC -> NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodeT
n
  Loop {} ->
    NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT -> m NodeT) -> NodeT -> m NodeT
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> [(NodeT, [EdgeT])] -> NodeT
DoNode Stm SOACS
s []
  Match {} ->
    NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT -> m NodeT) -> NodeT -> m NodeT
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> [(NodeT, [EdgeT])] -> NodeT
MatchNode Stm SOACS
s []
  Exp SOACS
e
    | [VName
output] <- Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
Pat (LetDec SOACS)
pat,
      Just (VName
ia, ArrayTransform
tr) <- Certs -> Exp SOACS -> Maybe (VName, ArrayTransform)
forall rep. Certs -> Exp rep -> Maybe (VName, ArrayTransform)
H.transformFromExp (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux) Exp SOACS
e ->
        NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT -> m NodeT) -> NodeT -> m NodeT
forall a b. (a -> b) -> a -> b
$ VName -> ArrayTransform -> VName -> NodeT
TransNode VName
output ArrayTransform
tr VName
ia
  Exp SOACS
_ -> NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodeT
n
nodeToSoacNode NodeT
n = NodeT -> m NodeT
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodeT
n

-- | Construct a graph with only nodes, but no edges.
emptyGraph :: Body SOACS -> DepGraph
emptyGraph :: Body SOACS -> DepGraph
emptyGraph Body SOACS
body =
  DepGraph
    { dgGraph :: Gr NodeT EdgeT
dgGraph = [DepNode] -> [DepEdge] -> Gr NodeT EdgeT
forall a b. [LNode a] -> [LEdge b] -> Gr a b
forall (gr :: * -> * -> *) a b.
Graph gr =>
[LNode a] -> [LEdge b] -> gr a b
G.mkGraph ([NodeT] -> [DepNode]
forall {b}. [b] -> [(Node, b)]
labelNodes ([NodeT]
stmnodes [NodeT] -> [NodeT] -> [NodeT]
forall a. Semigroup a => a -> a -> a
<> [NodeT]
resnodes [NodeT] -> [NodeT] -> [NodeT]
forall a. Semigroup a => a -> a -> a
<> [NodeT]
inputnodes)) [],
      dgProducerMapping :: ProducerMapping
dgProducerMapping = ProducerMapping
forall a. Monoid a => a
mempty,
      dgAliasTable :: AliasTable
dgAliasTable = AliasTable
aliases
    }
  where
    labelNodes :: [b] -> [(Node, b)]
labelNodes = [Node] -> [b] -> [(Node, b)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Node
0 ..]
    stmnodes :: [NodeT]
stmnodes = (Stm SOACS -> NodeT) -> [Stm SOACS] -> [NodeT]
forall a b. (a -> b) -> [a] -> [b]
map Stm SOACS -> NodeT
StmNode ([Stm SOACS] -> [NodeT]) -> [Stm SOACS] -> [NodeT]
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms SOACS -> [Stm SOACS]) -> Stms SOACS -> [Stm SOACS]
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Stms SOACS
forall rep. Body rep -> Stms rep
bodyStms Body SOACS
body
    resnodes :: [NodeT]
resnodes = (VName -> NodeT) -> [VName] -> [NodeT]
forall a b. (a -> b) -> [a] -> [b]
map VName -> NodeT
ResNode ([VName] -> [NodeT]) -> [VName] -> [NodeT]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Result -> Names
forall a. FreeIn a => a -> Names
freeIn (Result -> Names) -> Result -> Names
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Result
forall rep. Body rep -> Result
bodyResult Body SOACS
body
    inputnodes :: [NodeT]
inputnodes = (VName -> NodeT) -> [VName] -> [NodeT]
forall a b. (a -> b) -> [a] -> [b]
map VName -> NodeT
FreeNode ([VName] -> [NodeT]) -> [VName] -> [NodeT]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
consumed
    (Stms (Aliases SOACS)
_, (AliasTable
aliases, Names
consumed)) = AliasTable
-> Stms SOACS -> (Stms (Aliases SOACS), AliasesAndConsumed)
forall rep.
AliasableRep rep =>
AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
Alias.analyseStms AliasTable
forall a. Monoid a => a
mempty (Stms SOACS -> (Stms (Aliases SOACS), AliasesAndConsumed))
-> Stms SOACS -> (Stms (Aliases SOACS), AliasesAndConsumed)
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Stms SOACS
forall rep. Body rep -> Stms rep
bodyStms Body SOACS
body

getStmRes :: EdgeGenerator
getStmRes :: EdgeGenerator
getStmRes (ResNode VName
name) = [(VName
name, VName -> EdgeT
Res VName
name)]
getStmRes NodeT
_ = []

addResEdges :: (Monad m) => DepGraphAug m
addResEdges :: forall (m :: * -> *). Monad m => DepGraphAug m
addResEdges = EdgeGenerator -> DepGraphAug m
forall (m :: * -> *). Monad m => EdgeGenerator -> DepGraphAug m
augWithFun EdgeGenerator
getStmRes

-- | Make a dependency graph corresponding to a 'Body'.
mkDepGraph :: (HasScope SOACS m, Monad m) => Body SOACS -> m DepGraph
mkDepGraph :: forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Body SOACS -> m DepGraph
mkDepGraph Body SOACS
body = [DepGraphAug m] -> DepGraphAug m
forall (m :: * -> *). Monad m => [DepGraphAug m] -> DepGraphAug m
applyAugs [DepGraphAug m]
augs DepGraphAug m -> DepGraphAug m
forall a b. (a -> b) -> a -> b
$ Body SOACS -> DepGraph
emptyGraph Body SOACS
body
  where
    augs :: [DepGraphAug m]
augs =
      [ DepGraphAug m
forall (m :: * -> *). Monad m => DepGraphAug m
makeMapping,
        DepGraphAug m
forall (m :: * -> *). Monad m => DepGraphAug m
addDeps,
        DepGraphAug m
forall (m :: * -> *). Monad m => DepGraphAug m
addConsAndAliases,
        DepGraphAug m
forall (m :: * -> *). Monad m => DepGraphAug m
addExtraCons,
        DepGraphAug m
forall (m :: * -> *). Monad m => DepGraphAug m
addResEdges,
        (NodeT -> m NodeT) -> DepGraphAug m
forall (m :: * -> *).
Monad m =>
(NodeT -> m NodeT) -> DepGraphAug m
mapAcrossNodeTs NodeT -> m NodeT
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
NodeT -> m NodeT
nodeToSoacNode -- Must be done after adding edges
      ]

-- | Make a dependency graph corresponding to a function.
mkDepGraphForFun :: FunDef SOACS -> DepGraph
mkDepGraphForFun :: FunDef SOACS -> DepGraph
mkDepGraphForFun FunDef SOACS
f = Reader (Scope SOACS) DepGraph -> Scope SOACS -> DepGraph
forall r a. Reader r a -> r -> a
runReader (Body SOACS -> Reader (Scope SOACS) DepGraph
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Body SOACS -> m DepGraph
mkDepGraph (FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
f)) Scope SOACS
scope
  where
    scope :: Scope SOACS
scope = [Param (FParamInfo SOACS)] -> Scope SOACS
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (FunDef SOACS -> [Param (FParamInfo SOACS)]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef SOACS
f) 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 (Body SOACS -> Stms SOACS
forall rep. Body rep -> Stms rep
bodyStms (FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
f))

-- | Merges two contexts.
mergedContext :: (Ord b) => a -> G.Context a b -> G.Context a b -> G.Context a b
mergedContext :: forall b a. Ord b => a -> Context a b -> Context a b -> Context a b
mergedContext a
mergedlabel (Adj b
inp1, Node
n1, a
_, Adj b
out1) (Adj b
inp2, Node
n2, a
_, Adj b
out2) =
  let new_inp :: Adj b
new_inp = ((b, Node) -> Bool) -> Adj b -> Adj b
forall a. (a -> Bool) -> [a] -> [a]
filter (\(b, Node)
n -> (b, Node) -> Node
forall a b. (a, b) -> b
snd (b, Node)
n Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
/= Node
n1 Bool -> Bool -> Bool
&& (b, Node) -> Node
forall a b. (a, b) -> b
snd (b, Node)
n Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
/= Node
n2) (Adj b -> Adj b
forall a. Ord a => [a] -> [a]
nubOrd (Adj b
inp1 Adj b -> Adj b -> Adj b
forall a. Semigroup a => a -> a -> a
<> Adj b
inp2))
      new_out :: Adj b
new_out = ((b, Node) -> Bool) -> Adj b -> Adj b
forall a. (a -> Bool) -> [a] -> [a]
filter (\(b, Node)
n -> (b, Node) -> Node
forall a b. (a, b) -> b
snd (b, Node)
n Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
/= Node
n1 Bool -> Bool -> Bool
&& (b, Node) -> Node
forall a b. (a, b) -> b
snd (b, Node)
n Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
/= Node
n2) (Adj b -> Adj b
forall a. Ord a => [a] -> [a]
nubOrd (Adj b
out1 Adj b -> Adj b -> Adj b
forall a. Semigroup a => a -> a -> a
<> Adj b
out2))
   in (Adj b
new_inp, Node
n1, a
mergedlabel, Adj b
new_out)

-- | Remove the given node, and insert the 'DepContext' into the
-- graph, replacing any existing information about the node contained
-- in the 'DepContext'.
contractEdge :: (Monad m) => G.Node -> DepContext -> DepGraphAug m
contractEdge :: forall (m :: * -> *).
Monad m =>
Node -> DepContext -> DepGraphAug m
contractEdge Node
n2 DepContext
ctx DepGraph
dg = do
  let n1 :: Node
n1 = DepContext -> Node
forall a b. Context a b -> Node
G.node' DepContext
ctx -- n1 remains
  DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DepGraph -> m DepGraph) -> DepGraph -> m DepGraph
forall a b. (a -> b) -> a -> b
$ DepGraph
dg {dgGraph = ctx G.& G.delNodes [n1, n2] (dgGraph dg)}

-- Utils for fusibility/infusibility
-- find dependencies - either fusible or infusible. edges are generated based on these

-- | A classification of a free variable.
data Classification
  = -- | Used as array input to a SOAC (meaning fusible).
    SOACInput
  | -- | Used in some other way.
    Other
  deriving (Classification -> Classification -> Bool
(Classification -> Classification -> Bool)
-> (Classification -> Classification -> Bool) -> Eq Classification
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Classification -> Classification -> Bool
== :: Classification -> Classification -> Bool
$c/= :: Classification -> Classification -> Bool
/= :: Classification -> Classification -> Bool
Eq, Eq Classification
Eq Classification =>
(Classification -> Classification -> Ordering)
-> (Classification -> Classification -> Bool)
-> (Classification -> Classification -> Bool)
-> (Classification -> Classification -> Bool)
-> (Classification -> Classification -> Bool)
-> (Classification -> Classification -> Classification)
-> (Classification -> Classification -> Classification)
-> Ord Classification
Classification -> Classification -> Bool
Classification -> Classification -> Ordering
Classification -> Classification -> Classification
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Classification -> Classification -> Ordering
compare :: Classification -> Classification -> Ordering
$c< :: Classification -> Classification -> Bool
< :: Classification -> Classification -> Bool
$c<= :: Classification -> Classification -> Bool
<= :: Classification -> Classification -> Bool
$c> :: Classification -> Classification -> Bool
> :: Classification -> Classification -> Bool
$c>= :: Classification -> Classification -> Bool
>= :: Classification -> Classification -> Bool
$cmax :: Classification -> Classification -> Classification
max :: Classification -> Classification -> Classification
$cmin :: Classification -> Classification -> Classification
min :: Classification -> Classification -> Classification
Ord, Node -> Classification -> ShowS
[Classification] -> ShowS
Classification -> String
(Node -> Classification -> ShowS)
-> (Classification -> String)
-> ([Classification] -> ShowS)
-> Show Classification
forall a.
(Node -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Node -> Classification -> ShowS
showsPrec :: Node -> Classification -> ShowS
$cshow :: Classification -> String
show :: Classification -> String
$cshowList :: [Classification] -> ShowS
showList :: [Classification] -> ShowS
Show)

type Classifications = S.Set (VName, Classification)

freeClassifications :: (FreeIn a) => a -> Classifications
freeClassifications :: forall a. FreeIn a => a -> Classifications
freeClassifications =
  [(VName, Classification)] -> Classifications
forall a. Ord a => [a] -> Set a
S.fromList ([(VName, Classification)] -> Classifications)
-> (a -> [(VName, Classification)]) -> a -> Classifications
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([VName] -> [Classification] -> [(VName, Classification)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` Classification -> [Classification]
forall a. a -> [a]
repeat Classification
Other) ([VName] -> [(VName, Classification)])
-> (a -> [VName]) -> a -> [(VName, Classification)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList (Names -> [VName]) -> (a -> Names) -> a -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Names
forall a. FreeIn a => a -> Names
freeIn

stmInputs :: Stm SOACS -> Classifications
stmInputs :: Stm SOACS -> Classifications
stmInputs (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) =
  (Pat Type, StmAux ()) -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications (Pat Type
Pat (LetDec SOACS)
pat, StmAux ()
StmAux (ExpDec SOACS)
aux) Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> Exp SOACS -> Classifications
expInputs Exp SOACS
e

bodyInputs :: Body SOACS -> Classifications
bodyInputs :: Body SOACS -> Classifications
bodyInputs (Body BodyDec SOACS
_ Stms SOACS
stms Result
res) = (Stm SOACS -> Classifications) -> Stms SOACS -> Classifications
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm SOACS -> Classifications
stmInputs Stms SOACS
stms Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> Result -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications Result
res

expInputs :: Exp SOACS -> Classifications
expInputs :: Exp SOACS -> Classifications
expInputs (Match [SubExp]
cond [Case (Body SOACS)]
cases Body SOACS
defbody MatchDec (BranchType SOACS)
attr) =
  (Case (Body SOACS) -> Classifications)
-> [Case (Body SOACS)] -> Classifications
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Body SOACS -> Classifications
bodyInputs (Body SOACS -> Classifications)
-> (Case (Body SOACS) -> Body SOACS)
-> Case (Body SOACS)
-> Classifications
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (Body SOACS) -> Body SOACS
forall body. Case body -> body
caseBody) [Case (Body SOACS)]
cases
    Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> Body SOACS -> Classifications
bodyInputs Body SOACS
defbody
    Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> ([SubExp], MatchDec ExtType) -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications ([SubExp]
cond, MatchDec ExtType
MatchDec (BranchType SOACS)
attr)
expInputs (Loop [(Param (FParamInfo SOACS), SubExp)]
params LoopForm
form Body SOACS
b1) =
  ([(Param DeclType, SubExp)], LoopForm) -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications ([(Param DeclType, SubExp)]
[(Param (FParamInfo SOACS), SubExp)]
params, LoopForm
form) Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> Body SOACS -> Classifications
bodyInputs Body SOACS
b1
expInputs (Op Op SOACS
soac) = case Op SOACS
soac of
  Futhark.Screma SubExp
w [VName]
is ScremaForm SOACS
form -> [VName] -> Classifications
inputs [VName]
is Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> (SubExp, ScremaForm SOACS) -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications (SubExp
w, ScremaForm SOACS
form)
  Futhark.Hist SubExp
w [VName]
is [HistOp SOACS]
ops Lambda SOACS
lam -> [VName] -> Classifications
inputs [VName]
is Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> (SubExp, [HistOp SOACS], Lambda SOACS) -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications (SubExp
w, [HistOp SOACS]
ops, Lambda SOACS
lam)
  Futhark.Scatter SubExp
w [VName]
is ScatterSpec VName
lam Lambda SOACS
iws -> [VName] -> Classifications
inputs [VName]
is Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> (SubExp, ScatterSpec VName, Lambda SOACS) -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications (SubExp
w, ScatterSpec VName
lam, Lambda SOACS
iws)
  Futhark.Stream SubExp
w [VName]
is [SubExp]
nes Lambda SOACS
lam ->
    [VName] -> Classifications
inputs [VName]
is Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> (SubExp, [SubExp], Lambda SOACS) -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications (SubExp
w, [SubExp]
nes, Lambda SOACS
lam)
  Futhark.JVP {} -> SOAC SOACS -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications Op SOACS
SOAC SOACS
soac
  Futhark.VJP {} -> SOAC SOACS -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications Op SOACS
SOAC SOACS
soac
  where
    inputs :: [VName] -> Classifications
inputs = [(VName, Classification)] -> Classifications
forall a. Ord a => [a] -> Set a
S.fromList ([(VName, Classification)] -> Classifications)
-> ([VName] -> [(VName, Classification)])
-> [VName]
-> Classifications
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([VName] -> [Classification] -> [(VName, Classification)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` Classification -> [Classification]
forall a. a -> [a]
repeat Classification
SOACInput)
expInputs Exp SOACS
e
  | Just (VName
arr, ArrayTransform
_) <- Certs -> Exp SOACS -> Maybe (VName, ArrayTransform)
forall rep. Certs -> Exp rep -> Maybe (VName, ArrayTransform)
H.transformFromExp Certs
forall a. Monoid a => a
mempty Exp SOACS
e =
      (VName, Classification) -> Classifications
forall a. a -> Set a
S.singleton (VName
arr, Classification
SOACInput)
        Classifications -> Classifications -> Classifications
forall a. Semigroup a => a -> a -> a
<> Names -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications (Exp SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Exp SOACS
e Names -> Names -> Names
`namesSubtract` VName -> Names
oneName VName
arr)
  | Bool
otherwise = Exp SOACS -> Classifications
forall a. FreeIn a => a -> Classifications
freeClassifications Exp SOACS
e

stmNames :: Stm SOACS -> [VName]
stmNames :: Stm SOACS -> [VName]
stmNames = Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat Type -> [VName])
-> (Stm SOACS -> Pat Type) -> Stm SOACS -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm SOACS -> Pat Type
Stm SOACS -> Pat (LetDec SOACS)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat

getOutputs :: NodeT -> [VName]
getOutputs :: NodeT -> [VName]
getOutputs NodeT
node = case NodeT
node of
  (StmNode Stm SOACS
stm) -> Stm SOACS -> [VName]
stmNames Stm SOACS
stm
  (TransNode VName
v ArrayTransform
_ VName
_) -> [VName
v]
  (ResNode VName
_) -> []
  (FreeNode VName
name) -> [VName
name]
  (MatchNode Stm SOACS
stm [(NodeT, [EdgeT])]
_) -> Stm SOACS -> [VName]
stmNames Stm SOACS
stm
  (DoNode Stm SOACS
stm [(NodeT, [EdgeT])]
_) -> Stm SOACS -> [VName]
stmNames Stm SOACS
stm
  (SoacNode ArrayTransforms
_ Pat Type
pat SOAC SOACS
_ StmAux (ExpDec SOACS)
_) -> Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pat

-- | Is there a possibility of fusion?
isDep :: EdgeT -> Bool
isDep :: EdgeT -> Bool
isDep (Dep VName
_) = Bool
True
isDep (Res VName
_) = Bool
True
isDep EdgeT
_ = Bool
False

-- | Is this an infusible edge?
isInf :: (G.Node, G.Node, EdgeT) -> Bool
isInf :: DepEdge -> Bool
isInf (Node
_, Node
_, EdgeT
e) = case EdgeT
e of
  InfDep VName
_ -> Bool
True
  Fake VName
_ -> Bool
True -- this is infusible to avoid simultaneous cons/dep edges
  EdgeT
_ -> Bool
False

-- | Is this a 'Cons' edge?
isCons :: EdgeT -> Bool
isCons :: EdgeT -> Bool
isCons (Cons VName
_) = Bool
True
isCons EdgeT
_ = Bool
False

-- | Is this a withAcc?
isWithAccNodeT :: NodeT -> Bool
isWithAccNodeT :: NodeT -> Bool
isWithAccNodeT (StmNode (Let Pat (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ (WithAcc [WithAccInput SOACS]
_ Lambda SOACS
_))) = Bool
True
isWithAccNodeT NodeT
_ = Bool
False

isWithAccNodeId :: G.Node -> DepGraph -> Bool
isWithAccNodeId :: Node -> DepGraph -> Bool
isWithAccNodeId Node
node_id (DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
dgGraph = Gr NodeT EdgeT
g}) =
  let (Adj EdgeT
_, Node
_, NodeT
nT, Adj EdgeT
_) = Gr NodeT EdgeT -> Node -> DepContext
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Node -> Context a b
G.context Gr NodeT EdgeT
g Node
node_id
   in NodeT -> Bool
isWithAccNodeT NodeT
nT

unreachableEitherDir :: DepGraph -> G.Node -> G.Node -> Bool
unreachableEitherDir :: DepGraph -> Node -> Node -> Bool
unreachableEitherDir DepGraph
g Node
a Node
b =
  Bool -> Bool
not (DepGraph -> Node -> Node -> Bool
reachable DepGraph
g Node
a Node
b Bool -> Bool -> Bool
|| DepGraph -> Node -> Node -> Bool
reachable DepGraph
g Node
b Node
a)

vFusionFeasability :: DepGraph -> G.Node -> G.Node -> Bool
vFusionFeasability :: DepGraph -> Node -> Node -> Bool
vFusionFeasability dg :: DepGraph
dg@DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
dgGraph = Gr NodeT EdgeT
g} Node
n1 Node
n2 =
  (Node -> DepGraph -> Bool
isWithAccNodeId Node
n2 DepGraph
dg Bool -> Bool -> Bool
|| Bool -> Bool
not ((DepEdge -> Bool) -> [DepEdge] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any DepEdge -> Bool
isInf (DepGraph -> Node -> Node -> [DepEdge]
edgesBetween DepGraph
dg Node
n1 Node
n2)))
    Bool -> Bool -> Bool
&& Bool -> Bool
not ((Node -> Bool) -> [Node] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DepGraph -> Node -> Node -> Bool
reachable DepGraph
dg Node
n2) ((Node -> Bool) -> [Node] -> [Node]
forall a. (a -> Bool) -> [a] -> [a]
filter (Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
/= Node
n2) (Gr NodeT EdgeT -> Node -> [Node]
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Node -> [Node]
G.pre Gr NodeT EdgeT
g Node
n1)))

hFusionFeasability :: DepGraph -> G.Node -> G.Node -> Bool
hFusionFeasability :: DepGraph -> Node -> Node -> Bool
hFusionFeasability = DepGraph -> Node -> Node -> Bool
unreachableEitherDir