{-# LANGUAGE Strict #-}

-- | This module consists of rules for fusion
--     that involves WithAcc constructs.
--   Currently, we support two non-trivial
--   transformations:
--     I. SOAC-through-Trans-into-WithAcc fusion: a SoacNode is fused
--        into a WithAcc atomically when all dependency paths between
--        them go through TransNodes (reshapes/rearranges) that have no
--        other consumers. The canonical example is a map nest producing
--        multi-dimensional index and value arrays that are flattened
--        and consumed by a scatter. This must be done atomically to
--        avoid an infinite loop: absorbing only the TransNode causes
--        simplifyLambda to hoist the cheap reshape back out, recreating
--        the same TransNode and triggering the rule again. The strategy
--        is to prepend the SoacNode and all TransNode statements into
--        the WithAcc lambda body and run doFusionInLambda to fuse
--        further where possible.
--
--    II. WithAcc-WithAcc fusion: two withaccs can be
--        fused as long as the common accumulators use
--        the same operator, and as long as the non-accumulator
--        input of an WithAcc is not used as an accumulator in
--        the other. This fusion opens the door for fusing
--        the SOACs appearing inside the WithAccs. This is
--        also intended to demonstrate that it is not so
--        important where exactly the WithAccs were originally
--        introduced in the code, it is more important that
--        they can be transformed by various optimizations passes.
module Futhark.Optimise.Fusion.RulesWithAccs
  ( tryFuseWithAccs,
    trySoacThroughTransIntoWithAcc,
  )
where

import Control.Monad
import Data.Graph.Inductive.Graph qualified as G
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe (maybeToList)
import Futhark.Analysis.HORep.SOAC qualified as H
import Futhark.Construct
import Futhark.IR.SOACS hiding (SOAC (..))
import Futhark.Optimise.Fusion.GraphRep
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util (nubOrd)

---------------------------------------------------
--- II. WithAcc-WithAcc Fusion
---------------------------------------------------

-- | Local helper type that tuples together:
--   1.   the pattern element corresponding to one withacc input
--   2.   the withacc input
--   3-5  withacc's lambda corresponding acc-certificate param,
--           argument param and result name
data AccTup = AccTup
  { AccTup -> [PatElem (LetDec SOACS)]
accTup1 :: [PatElem (LetDec SOACS)],
    AccTup -> WithAccInput SOACS
accTup2 :: WithAccInput SOACS,
    AccTup -> LParam SOACS
accTup3 :: LParam SOACS,
    AccTup -> LParam SOACS
accTup4 :: LParam SOACS,
    AccTup -> (VName, Certs)
accTup5 :: (VName, Certs)
  }
  deriving (Int -> AccTup -> ShowS
[AccTup] -> ShowS
AccTup -> String
(Int -> AccTup -> ShowS)
-> (AccTup -> String) -> ([AccTup] -> ShowS) -> Show AccTup
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AccTup -> ShowS
showsPrec :: Int -> AccTup -> ShowS
$cshow :: AccTup -> String
show :: AccTup -> String
$cshowList :: [AccTup] -> ShowS
showList :: [AccTup] -> ShowS
Show)

groupAccs ::
  [PatElem (LetDec SOACS)] ->
  [WithAccInput SOACS] ->
  Lambda SOACS ->
  ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccs :: [PatElem (LetDec SOACS)]
-> [WithAccInput SOACS]
-> Lambda SOACS
-> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccs [PatElem (LetDec SOACS)]
pat_els [WithAccInput SOACS]
wacc_inps Lambda SOACS
wlam =
  let lam_params :: [LParam SOACS]
lam_params = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
wlam
      n :: Int
n = [Param Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param Type]
lam_params
      ([Param Type]
lam_par_crts, [Param Type]
lam_par_accs) = Int -> [Param Type] -> ([Param Type], [Param Type])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [Param Type]
lam_params
      lab_res_ses :: [SubExpRes]
lab_res_ses = GBody SOACS SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody SOACS SubExpRes -> [SubExpRes])
-> GBody SOACS SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
wlam
   in [PatElem (LetDec SOACS)]
-> [WithAccInput SOACS]
-> [LParam SOACS]
-> [LParam SOACS]
-> [SubExpRes]
-> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccsHlp [PatElem (LetDec SOACS)]
pat_els [WithAccInput SOACS]
wacc_inps [Param Type]
[LParam SOACS]
lam_par_crts [Param Type]
[LParam SOACS]
lam_par_accs [SubExpRes]
lab_res_ses

groupAccsHlp ::
  [PatElem (LetDec SOACS)] ->
  [WithAccInput SOACS] ->
  [LParam SOACS] ->
  [LParam SOACS] ->
  [SubExpRes] ->
  ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccsHlp :: [PatElem (LetDec SOACS)]
-> [WithAccInput SOACS]
-> [LParam SOACS]
-> [LParam SOACS]
-> [SubExpRes]
-> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccsHlp [PatElem (LetDec SOACS)]
pat_els [] [] [] [SubExpRes]
lam_res_ses
  | [PatElem Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PatElem Type]
[PatElem (LetDec SOACS)]
pat_els Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [SubExpRes] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExpRes]
lam_res_ses =
      ([], [PatElem Type] -> [SubExpRes] -> [(PatElem Type, SubExpRes)]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatElem Type]
[PatElem (LetDec SOACS)]
pat_els [SubExpRes]
lam_res_ses)
groupAccsHlp
  [PatElem (LetDec SOACS)]
pat_els
  (winp :: WithAccInput SOACS
winp@(ShapeBase SubExp
_, [VName]
inp, Maybe (Lambda SOACS, [SubExp])
_) : [WithAccInput SOACS]
wacc_inps)
  (LParam SOACS
par_crt : [LParam SOACS]
lam_par_crts)
  (LParam SOACS
par_acc : [LParam SOACS]
lam_par_accs)
  (SubExpRes
res_se : [SubExpRes]
lam_res_ses)
    | Int
n <- [VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
inp,
      Var VName
res_nm <- SubExpRes -> SubExp
resSubExp SubExpRes
res_se =
        let ([PatElem Type]
pat_els_cur, [PatElem Type]
pat_els') = Int -> [PatElem Type] -> ([PatElem Type], [PatElem Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [PatElem Type]
[PatElem (LetDec SOACS)]
pat_els
            ([AccTup]
rec1, [(PatElem (LetDec SOACS), SubExpRes)]
rec2) = [PatElem (LetDec SOACS)]
-> [WithAccInput SOACS]
-> [LParam SOACS]
-> [LParam SOACS]
-> [SubExpRes]
-> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccsHlp [PatElem Type]
[PatElem (LetDec SOACS)]
pat_els' [WithAccInput SOACS]
wacc_inps [LParam SOACS]
lam_par_crts [LParam SOACS]
lam_par_accs [SubExpRes]
lam_res_ses
         in ([PatElem (LetDec SOACS)]
-> WithAccInput SOACS
-> LParam SOACS
-> LParam SOACS
-> (VName, Certs)
-> AccTup
AccTup [PatElem Type]
[PatElem (LetDec SOACS)]
pat_els_cur WithAccInput SOACS
winp LParam SOACS
par_crt LParam SOACS
par_acc (VName
res_nm, SubExpRes -> Certs
resCerts SubExpRes
res_se) AccTup -> [AccTup] -> [AccTup]
forall a. a -> [a] -> [a]
: [AccTup]
rec1, [(PatElem Type, SubExpRes)]
[(PatElem (LetDec SOACS), SubExpRes)]
rec2)
groupAccsHlp [PatElem (LetDec SOACS)]
_ [WithAccInput SOACS]
_ [LParam SOACS]
_ [LParam SOACS]
_ [SubExpRes]
_ =
  String -> ([AccTup], [(PatElem Type, SubExpRes)])
forall a. HasCallStack => String -> a
error String
"Unreachable case reached in groupAccsHlp!"

matchingAccTup :: AccTup -> AccTup -> Bool
matchingAccTup :: AccTup -> AccTup -> Bool
matchingAccTup
  (AccTup [PatElem (LetDec SOACS)]
pat_els1 (ShapeBase SubExp
shp1, [VName]
_winp_arrs1, Maybe (Lambda SOACS, [SubExp])
mlam1) LParam SOACS
_ LParam SOACS
_ (VName, Certs)
_)
  (AccTup [PatElem (LetDec SOACS)]
_ (ShapeBase SubExp
shp2, [VName]
winp_arrs2, Maybe (Lambda SOACS, [SubExp])
mlam2) LParam SOACS
_ LParam SOACS
_ (VName, Certs)
_) =
    ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
shp1 [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
shp2
      Bool -> Bool -> Bool
&& (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]
[PatElem (LetDec SOACS)]
pat_els1 [VName] -> [VName] -> Bool
forall a. Eq a => a -> a -> Bool
== [VName]
winp_arrs2
      Bool -> Bool -> Bool
&& case (Maybe (Lambda SOACS, [SubExp])
mlam1, Maybe (Lambda SOACS, [SubExp])
mlam2) of
        (Maybe (Lambda SOACS, [SubExp])
Nothing, Maybe (Lambda SOACS, [SubExp])
Nothing) -> Bool
True
        (Just (Lambda SOACS
lam1, [SubExp]
see1), Just (Lambda SOACS
lam2, [SubExp]
see2)) ->
          [SubExp]
see1 [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== [SubExp]
see2 Bool -> Bool -> Bool
&& Map VName VName -> Lambda SOACS -> Lambda SOACS -> Bool
equivLambda Map VName VName
forall k a. Map k a
M.empty Lambda SOACS
lam1 Lambda SOACS
lam2
        (Maybe (Lambda SOACS, [SubExp]), Maybe (Lambda SOACS, [SubExp]))
_ -> Bool
False

groupCommonAccs :: [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
groupCommonAccs :: [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
groupCommonAccs [] [AccTup]
tup_accs2 =
  ([], [], [AccTup]
tup_accs2)
groupCommonAccs (AccTup
tup_acc1 : [AccTup]
tup_accs1) [AccTup]
tup_accs2
  | ([AccTup]
commons2, [AccTup]
uncommons2) <- (AccTup -> Bool) -> [AccTup] -> ([AccTup], [AccTup])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (AccTup -> AccTup -> Bool
matchingAccTup AccTup
tup_acc1) [AccTup]
tup_accs2,
    [AccTup] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [AccTup]
commons2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 =
      let ([(AccTup, AccTup)]
rec1, [AccTup]
rec2, [AccTup]
rec3) = [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
groupCommonAccs [AccTup]
tup_accs1 [AccTup]
uncommons2
       in if [AccTup] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [AccTup]
commons2
            then ([(AccTup, AccTup)]
rec1, AccTup
tup_acc1 AccTup -> [AccTup] -> [AccTup]
forall a. a -> [a] -> [a]
: [AccTup]
rec2, [AccTup]
rec3)
            else ((AccTup
tup_acc1, [AccTup] -> AccTup
forall a. HasCallStack => [a] -> a
head [AccTup]
commons2) (AccTup, AccTup) -> [(AccTup, AccTup)] -> [(AccTup, AccTup)]
forall a. a -> [a] -> [a]
: [(AccTup, AccTup)]
rec1, [AccTup]
rec2, [AccTup]
rec3)
groupCommonAccs [AccTup]
_ [AccTup]
_ =
  String -> ([(AccTup, AccTup)], [AccTup], [AccTup])
forall a. HasCallStack => String -> a
error String
"Unreachable case reached in groupCommonAccs!"

-- | Simple case for fusing two withAccs (can be extended):
--    let (b1, ..., bm, x1, ..., xq) = withAcc a1 ... am lam1
--    let (d1, ..., dn, y1, ..., yp) = withAcc c1 ... cn lam2
-- Notation: `b1 ... bm` are the accumulator results of the
--     first withAcc and `d1, ..., dn` of the second withAcc.
--     `x1 ... xq` and `y1, ..., yp` are non-accumulator results.
-- Conservative conditions:
--   1. for any bi (i=1..m) either `bi IN {c1, ..., cm}` OR
--        `bi NOT-IN FV(lam2)`, i.e., perfect producer-consumer
--        relation on accums. Of course the binary-op should
--        be the same.
--   2. The `bs` that are also accumulated upon in lam2
--        do NOT belong to the `infusible` set (they are destroyed)
--   3. x1 ... xq do not overlap with c1 ... cn
-- Fusion will create one withacc that accumulates on the
--   union of `a1 ... am` and `c1 ... cn` and returns, in addition
--   to the accumulator arrays the union of regular variables
--   `x1 ... xq` and `y1, ..., yp`
tryFuseWithAccs ::
  (HasScope SOACS m, MonadFreshNames m) =>
  [VName] ->
  Stm SOACS ->
  Stm SOACS ->
  Maybe (m (Stm SOACS))
tryFuseWithAccs :: forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
[VName] -> Stm SOACS -> Stm SOACS -> Maybe (m (Stm SOACS))
tryFuseWithAccs
  [VName]
infusible
  (Let Pat (LetDec SOACS)
pat1 StmAux (ExpDec SOACS)
aux1 (WithAcc [WithAccInput SOACS]
w_inps1 Lambda SOACS
lam1))
  (Let Pat (LetDec SOACS)
pat2 StmAux (ExpDec SOACS)
aux2 (WithAcc [WithAccInput SOACS]
w_inps2 Lambda SOACS
lam2))
    | ([PatElem Type]
pat1_els, [PatElem Type]
pat2_els) <- (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat1, Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat2),
      ([AccTup]
acc_tup1, [(PatElem (LetDec SOACS), SubExpRes)]
other_pr1) <- [PatElem (LetDec SOACS)]
-> [WithAccInput SOACS]
-> Lambda SOACS
-> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccs [PatElem Type]
[PatElem (LetDec SOACS)]
pat1_els [WithAccInput SOACS]
w_inps1 Lambda SOACS
lam1,
      ([AccTup]
acc_tup2, [(PatElem (LetDec SOACS), SubExpRes)]
other_pr2) <- [PatElem (LetDec SOACS)]
-> [WithAccInput SOACS]
-> Lambda SOACS
-> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccs [PatElem Type]
[PatElem (LetDec SOACS)]
pat2_els [WithAccInput SOACS]
w_inps2 Lambda SOACS
lam2,
      ([(AccTup, AccTup)]
tup_common, [AccTup]
acc_tup1', [AccTup]
acc_tup2') <-
        [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
groupCommonAccs [AccTup]
acc_tup1 [AccTup]
acc_tup2,
      -- safety 0: make sure that the accs from acc_tup1' and
      --           acc_tup2' do not overlap
      [VName]
pnms_1' <- (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
$ (AccTup -> [PatElem Type]) -> [AccTup] -> [PatElem Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap AccTup -> [PatElem Type]
AccTup -> [PatElem (LetDec SOACS)]
accTup1 [AccTup]
acc_tup1',
      [VName]
winp_2' <- (AccTup -> [VName]) -> [AccTup] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((\(ShapeBase SubExp
_, [VName]
xs, Maybe (Lambda SOACS, [SubExp])
_) -> [VName]
xs) (WithAccInput SOACS -> [VName])
-> (AccTup -> WithAccInput SOACS) -> AccTup -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AccTup -> WithAccInput SOACS
accTup2) [AccTup]
acc_tup2',
      Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Names -> Names -> Bool
namesIntersect ([VName] -> Names
namesFromList [VName]
pnms_1') ([VName] -> Names
namesFromList [VName]
winp_2'),
      -- safety 1: we have already determined the commons;
      --           now we also need to check NOT-IN FV(lam2)
      Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Names -> Names -> Bool
namesIntersect ([VName] -> Names
namesFromList [VName]
pnms_1') (Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
lam2),
      -- safety 2:
      [VName]
bs <- (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
$ ((AccTup, AccTup) -> [PatElem Type])
-> [(AccTup, AccTup)] -> [PatElem Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (AccTup -> [PatElem Type]
AccTup -> [PatElem (LetDec SOACS)]
accTup1 (AccTup -> [PatElem Type])
-> ((AccTup, AccTup) -> AccTup)
-> (AccTup, AccTup)
-> [PatElem Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AccTup, AccTup) -> AccTup
forall a b. (a, b) -> a
fst) [(AccTup, AccTup)]
tup_common,
      (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
infusible) [VName]
bs,
      -- safety 3:
      Names
cs <- [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (AccTup -> [VName]) -> [AccTup] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((\(ShapeBase SubExp
_, [VName]
xs, Maybe (Lambda SOACS, [SubExp])
_) -> [VName]
xs) (WithAccInput SOACS -> [VName])
-> (AccTup -> WithAccInput SOACS) -> AccTup -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AccTup -> WithAccInput SOACS
accTup2) [AccTup]
acc_tup2,
      ((PatElem Type, SubExpRes) -> Bool)
-> [(PatElem Type, SubExpRes)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((VName -> Names -> Bool
`notNameIn` Names
cs) (VName -> Bool)
-> ((PatElem Type, SubExpRes) -> VName)
-> (PatElem Type, SubExpRes)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName (PatElem Type -> VName)
-> ((PatElem Type, SubExpRes) -> PatElem Type)
-> (PatElem Type, SubExpRes)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PatElem Type, SubExpRes) -> PatElem Type
forall a b. (a, b) -> a
fst) [(PatElem Type, SubExpRes)]
[(PatElem (LetDec SOACS), SubExpRes)]
other_pr1 = m (Stm SOACS) -> Maybe (m (Stm SOACS))
forall a. a -> Maybe a
Just (m (Stm SOACS) -> Maybe (m (Stm SOACS)))
-> m (Stm SOACS) -> Maybe (m (Stm SOACS))
forall a b. (a -> b) -> a -> b
$ do
        let getCertPairs :: (AccTup, AccTup) -> (VName, VName)
getCertPairs (AccTup
t1, AccTup
t2) = (Param Type -> VName
forall dec. Param dec -> VName
paramName (AccTup -> LParam SOACS
accTup3 AccTup
t2), Param Type -> VName
forall dec. Param dec -> VName
paramName (AccTup -> LParam SOACS
accTup3 AccTup
t1))
            tab_certs :: Map VName VName
tab_certs = [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$ ((AccTup, AccTup) -> (VName, VName))
-> [(AccTup, AccTup)] -> [(VName, VName)]
forall a b. (a -> b) -> [a] -> [b]
map (AccTup, AccTup) -> (VName, VName)
getCertPairs [(AccTup, AccTup)]
tup_common
            lam2_bdy' :: GBody SOACS SubExpRes
lam2_bdy' = Map VName VName -> GBody SOACS SubExpRes -> GBody SOACS SubExpRes
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
tab_certs (Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam2)
            rcrt_params :: [Param Type]
rcrt_params = ((AccTup, AccTup) -> Param Type)
-> [(AccTup, AccTup)] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (AccTup -> Param Type
AccTup -> LParam SOACS
accTup3 (AccTup -> Param Type)
-> ((AccTup, AccTup) -> AccTup) -> (AccTup, AccTup) -> Param Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AccTup, AccTup) -> AccTup
forall a b. (a, b) -> a
fst) [(AccTup, AccTup)]
tup_common [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ (AccTup -> Param Type) -> [AccTup] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map AccTup -> Param Type
AccTup -> LParam SOACS
accTup3 [AccTup]
acc_tup1' [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ (AccTup -> Param Type) -> [AccTup] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map AccTup -> Param Type
AccTup -> LParam SOACS
accTup3 [AccTup]
acc_tup2'
            racc_params :: [Param Type]
racc_params = ((AccTup, AccTup) -> Param Type)
-> [(AccTup, AccTup)] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (AccTup -> Param Type
AccTup -> LParam SOACS
accTup4 (AccTup -> Param Type)
-> ((AccTup, AccTup) -> AccTup) -> (AccTup, AccTup) -> Param Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AccTup, AccTup) -> AccTup
forall a b. (a, b) -> a
fst) [(AccTup, AccTup)]
tup_common [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ (AccTup -> Param Type) -> [AccTup] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map AccTup -> Param Type
AccTup -> LParam SOACS
accTup4 [AccTup]
acc_tup1' [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ (AccTup -> Param Type) -> [AccTup] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map AccTup -> Param Type
AccTup -> LParam SOACS
accTup4 [AccTup]
acc_tup2'
            ([VName]
comm_res_nms, [Certs]
comm_res_certs2) = [(VName, Certs)] -> ([VName], [Certs])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, Certs)] -> ([VName], [Certs]))
-> [(VName, Certs)] -> ([VName], [Certs])
forall a b. (a -> b) -> a -> b
$ ((AccTup, AccTup) -> (VName, Certs))
-> [(AccTup, AccTup)] -> [(VName, Certs)]
forall a b. (a -> b) -> [a] -> [b]
map (AccTup -> (VName, Certs)
accTup5 (AccTup -> (VName, Certs))
-> ((AccTup, AccTup) -> AccTup)
-> (AccTup, AccTup)
-> (VName, Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AccTup, AccTup) -> AccTup
forall a b. (a, b) -> b
snd) [(AccTup, AccTup)]
tup_common
            ([VName]
_, [Certs]
comm_res_certs1) = [(VName, Certs)] -> ([VName], [Certs])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, Certs)] -> ([VName], [Certs]))
-> [(VName, Certs)] -> ([VName], [Certs])
forall a b. (a -> b) -> a -> b
$ ((AccTup, AccTup) -> (VName, Certs))
-> [(AccTup, AccTup)] -> [(VName, Certs)]
forall a b. (a -> b) -> [a] -> [b]
map (AccTup -> (VName, Certs)
accTup5 (AccTup -> (VName, Certs))
-> ((AccTup, AccTup) -> AccTup)
-> (AccTup, AccTup)
-> (VName, Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AccTup, AccTup) -> AccTup
forall a b. (a, b) -> a
fst) [(AccTup, AccTup)]
tup_common
            com_res_certs :: [Certs]
com_res_certs = (Certs -> Certs -> Certs) -> [Certs] -> [Certs] -> [Certs]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Certs
x Certs
y -> [VName] -> Certs
Certs (Certs -> [VName]
unCerts Certs
x [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ Certs -> [VName]
unCerts Certs
y)) [Certs]
comm_res_certs1 [Certs]
comm_res_certs2
            bdyres_certs :: [Certs]
bdyres_certs = [Certs]
com_res_certs [Certs] -> [Certs] -> [Certs]
forall a. [a] -> [a] -> [a]
++ (AccTup -> Certs) -> [AccTup] -> [Certs]
forall a b. (a -> b) -> [a] -> [b]
map ((VName, Certs) -> Certs
forall a b. (a, b) -> b
snd ((VName, Certs) -> Certs)
-> (AccTup -> (VName, Certs)) -> AccTup -> Certs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AccTup -> (VName, Certs)
accTup5) ([AccTup]
acc_tup1' [AccTup] -> [AccTup] -> [AccTup]
forall a. [a] -> [a] -> [a]
++ [AccTup]
acc_tup2')
            bdyres_accse :: [SubExp]
bdyres_accse = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
comm_res_nms [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ (AccTup -> SubExp) -> [AccTup] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp) -> (AccTup -> VName) -> AccTup -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Certs) -> VName
forall a b. (a, b) -> a
fst ((VName, Certs) -> VName)
-> (AccTup -> (VName, Certs)) -> AccTup -> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AccTup -> (VName, Certs)
accTup5) ([AccTup]
acc_tup1' [AccTup] -> [AccTup] -> [AccTup]
forall a. [a] -> [a] -> [a]
++ [AccTup]
acc_tup2')
            bdy_res_accs :: [SubExpRes]
bdy_res_accs = (Certs -> SubExp -> SubExpRes)
-> [Certs] -> [SubExp] -> [SubExpRes]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Certs -> SubExp -> SubExpRes
SubExpRes [Certs]
bdyres_certs [SubExp]
bdyres_accse
            bdy_res_others :: [SubExpRes]
bdy_res_others = ((PatElem Type, SubExpRes) -> SubExpRes)
-> [(PatElem Type, SubExpRes)] -> [SubExpRes]
forall a b. (a -> b) -> [a] -> [b]
map (PatElem Type, SubExpRes) -> SubExpRes
forall a b. (a, b) -> b
snd ([(PatElem Type, SubExpRes)] -> [SubExpRes])
-> [(PatElem Type, SubExpRes)] -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ [(PatElem Type, SubExpRes)]
[(PatElem (LetDec SOACS), SubExpRes)]
other_pr1 [(PatElem Type, SubExpRes)]
-> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
forall a. [a] -> [a] -> [a]
++ [(PatElem Type, SubExpRes)]
[(PatElem (LetDec SOACS), SubExpRes)]
other_pr2
        res_lam <-
          [LParam SOACS] -> Builder SOACS [SubExpRes] -> m (Lambda SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
[LParam rep] -> Builder rep [SubExpRes] -> m (Lambda rep)
runLambdaBuilder ([Param Type]
rcrt_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
racc_params) (Builder SOACS [SubExpRes] -> m (Lambda SOACS))
-> Builder SOACS [SubExpRes] -> m (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ do
            -- add the stms of lam1
            (Stm SOACS -> BuilderT SOACS (State VNameSource) ())
-> [Stm SOACS] -> BuilderT SOACS (State VNameSource) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
Stm SOACS -> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm ([Stm SOACS] -> BuilderT SOACS (State VNameSource) ())
-> [Stm SOACS] -> BuilderT SOACS (State VNameSource) ()
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
$ GBody SOACS SubExpRes -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody SOACS SubExpRes -> Stms SOACS)
-> GBody SOACS SubExpRes -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam1
            -- add the copy stms for the common accumulator
            [(AccTup, AccTup)]
-> ((AccTup, AccTup) -> BuilderT SOACS (State VNameSource) ())
-> BuilderT SOACS (State VNameSource) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(AccTup, AccTup)]
tup_common (((AccTup, AccTup) -> BuilderT SOACS (State VNameSource) ())
 -> BuilderT SOACS (State VNameSource) ())
-> ((AccTup, AccTup) -> BuilderT SOACS (State VNameSource) ())
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ \(AccTup
tup1, AccTup
tup2) -> do
              let (Param Type
lpar1, Param Type
lpar2) = (AccTup -> LParam SOACS
accTup4 AccTup
tup1, AccTup -> LParam SOACS
accTup4 AccTup
tup2)
                  ((VName
nm1, Certs
_), VName
nm2, Type
tp_acc) = (AccTup -> (VName, Certs)
accTup5 AccTup
tup1, Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
lpar2, Param Type -> Type
forall dec. Param dec -> dec
paramDec Param Type
lpar1)
              Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind ([PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
nm2 Type
tp_acc]) (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) ())
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> 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 -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
nm1
            -- add copy stms to bring in scope x1 ... xq
            [(PatElem Type, SubExpRes)]
-> ((PatElem Type, SubExpRes)
    -> BuilderT SOACS (State VNameSource) ())
-> BuilderT SOACS (State VNameSource) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(PatElem Type, SubExpRes)]
[(PatElem (LetDec SOACS), SubExpRes)]
other_pr1 (((PatElem Type, SubExpRes)
  -> BuilderT SOACS (State VNameSource) ())
 -> BuilderT SOACS (State VNameSource) ())
-> ((PatElem Type, SubExpRes)
    -> BuilderT SOACS (State VNameSource) ())
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ \(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)
              Certs
-> BuilderT SOACS (State VNameSource) ()
-> BuilderT SOACS (State VNameSource) ()
forall a.
Certs
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying (SubExpRes -> Certs
resCerts SubExpRes
bdy_res) (BuilderT SOACS (State VNameSource) ()
 -> BuilderT SOACS (State VNameSource) ())
-> BuilderT SOACS (State VNameSource) ()
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
                Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind ([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]) (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) ())
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall rep. BasicOp -> Exp rep
BasicOp (SubExp -> BasicOp
SubExp SubExp
se)
            -- add the statements of lam2 (in which the acc-certificates have been substituted)
            (Stm SOACS -> BuilderT SOACS (State VNameSource) ())
-> [Stm SOACS] -> BuilderT SOACS (State VNameSource) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
Stm SOACS -> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm ([Stm SOACS] -> BuilderT SOACS (State VNameSource) ())
-> [Stm SOACS] -> BuilderT SOACS (State VNameSource) ()
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
$ GBody SOACS SubExpRes -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms GBody SOACS SubExpRes
lam2_bdy'
            -- build the result of body
            [SubExpRes] -> Builder SOACS [SubExpRes]
forall a. a -> BuilderT SOACS (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExpRes] -> Builder SOACS [SubExpRes])
-> [SubExpRes] -> Builder SOACS [SubExpRes]
forall a b. (a -> b) -> a -> b
$ [SubExpRes]
bdy_res_accs [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. [a] -> [a] -> [a]
++ [SubExpRes]
bdy_res_others
        res_lam' <- renameLambda res_lam
        let res_pat =
              ((AccTup, AccTup) -> [PatElem Type])
-> [(AccTup, AccTup)] -> [PatElem Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (AccTup -> [PatElem Type]
AccTup -> [PatElem (LetDec SOACS)]
accTup1 (AccTup -> [PatElem Type])
-> ((AccTup, AccTup) -> AccTup)
-> (AccTup, AccTup)
-> [PatElem Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AccTup, AccTup) -> AccTup
forall a b. (a, b) -> b
snd) [(AccTup, AccTup)]
tup_common
                [PatElem Type] -> [PatElem Type] -> [PatElem Type]
forall a. [a] -> [a] -> [a]
++ (AccTup -> [PatElem Type]) -> [AccTup] -> [PatElem Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap AccTup -> [PatElem Type]
AccTup -> [PatElem (LetDec SOACS)]
accTup1 ([AccTup]
acc_tup1' [AccTup] -> [AccTup] -> [AccTup]
forall a. [a] -> [a] -> [a]
++ [AccTup]
acc_tup2')
                [PatElem Type] -> [PatElem Type] -> [PatElem Type]
forall a. [a] -> [a] -> [a]
++ ((PatElem Type, SubExpRes) -> PatElem Type)
-> [(PatElem Type, SubExpRes)] -> [PatElem Type]
forall a b. (a -> b) -> [a] -> [b]
map (PatElem Type, SubExpRes) -> PatElem Type
forall a b. (a, b) -> a
fst ([(PatElem Type, SubExpRes)]
[(PatElem (LetDec SOACS), SubExpRes)]
other_pr1 [(PatElem Type, SubExpRes)]
-> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
forall a. [a] -> [a] -> [a]
++ [(PatElem Type, SubExpRes)]
[(PatElem (LetDec SOACS), SubExpRes)]
other_pr2)
            res_w_inps = ((AccTup, AccTup) -> WithAccInput SOACS)
-> [(AccTup, AccTup)] -> [WithAccInput SOACS]
forall a b. (a -> b) -> [a] -> [b]
map (AccTup -> WithAccInput SOACS
accTup2 (AccTup -> WithAccInput SOACS)
-> ((AccTup, AccTup) -> AccTup)
-> (AccTup, AccTup)
-> WithAccInput SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AccTup, AccTup) -> AccTup
forall a b. (a, b) -> a
fst) [(AccTup, AccTup)]
tup_common [WithAccInput SOACS]
-> [WithAccInput SOACS] -> [WithAccInput SOACS]
forall a. [a] -> [a] -> [a]
++ (AccTup -> WithAccInput SOACS) -> [AccTup] -> [WithAccInput SOACS]
forall a b. (a -> b) -> [a] -> [b]
map AccTup -> WithAccInput SOACS
accTup2 ([AccTup]
acc_tup1' [AccTup] -> [AccTup] -> [AccTup]
forall a. [a] -> [a] -> [a]
++ [AccTup]
acc_tup2')
        res_w_inps' <- mapM renameLamInWAccInp res_w_inps
        pure $ Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam'
    where
      renameLamInWAccInp :: (a, b, Maybe (GLambda rep t, b))
-> m (a, b, Maybe (GLambda rep t, b))
renameLamInWAccInp (a
shp, b
inps, Just (GLambda rep t
lam, b
se)) = do
        lam' <- GLambda rep t -> m (GLambda rep t)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda GLambda rep t
lam
        pure (shp, inps, Just (lam', se))
      renameLamInWAccInp (a, b, Maybe (GLambda rep t, b))
winp = (a, b, Maybe (GLambda rep t, b))
-> m (a, b, Maybe (GLambda rep t, b))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a, b, Maybe (GLambda rep t, b))
winp
--
tryFuseWithAccs [VName]
_ Stm SOACS
_ Stm SOACS
_ =
  Maybe (m (Stm SOACS))
forall a. Maybe a
Nothing

---------------------------------------------------
--- I. SOAC-through-Trans-into-WithAcc Fusion
---------------------------------------------------

-- | See the module-level description of transformation I.
-- The @doFusionInLambda@ and @fusedSomething@ arguments are callbacks
-- from Fusion.hs to avoid a circular import.
trySoacThroughTransIntoWithAcc ::
  (HasScope SOACS m, MonadFreshNames m) =>
  (Lambda SOACS -> m (Lambda SOACS, Bool)) ->
  (NodeT -> m (Maybe NodeT)) ->
  G.Node ->
  DepGraphAug m
trySoacThroughTransIntoWithAcc :: forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
(Lambda SOACS -> m (Lambda SOACS, Bool))
-> (NodeT -> m (Maybe NodeT)) -> Int -> DepGraphAug m
trySoacThroughTransIntoWithAcc Lambda SOACS -> m (Lambda SOACS, Bool)
doFusionInLambda NodeT -> m (Maybe NodeT)
fusedSomething Int
wacc_id 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
wacc_id Gr NodeT EdgeT
g) = DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  | Just (StmNode (Let Pat (LetDec SOACS)
pat2 StmAux (ExpDec SOACS)
aux2 (WithAcc [WithAccInput SOACS]
w_inps Lambda SOACS
lam0))) <- Gr NodeT EdgeT -> Int -> Maybe NodeT
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Maybe a
G.lab Gr NodeT EdgeT
g Int
wacc_id = do
      -- Edges go FROM consumers TO producers:
      --   G.lpre g n = consumers of n; G.lsuc g n = producers n depends on.
      -- realConsumers n: consumers of n, excluding Alias edges and self-loops.
      let wacc_cons_nms :: 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
          realConsumers :: Int -> [Int]
realConsumers Int
n =
            [Int] -> [Int]
forall a. Ord a => [a] -> [a]
nubOrd ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$
              ((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
m, EdgeT
e) -> Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n Bool -> Bool -> Bool
&& case EdgeT
e of Alias {} -> Bool
False; EdgeT
_ -> Bool
True) ([(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
n
          -- TransNodes that directly feed wacc and are exclusively consumed by wacc.
          trans_preds :: [(Int, VName, ArrayTransform, VName)]
trans_preds = do
            (tn_id, _) <- Gr NodeT EdgeT -> Int -> [(Int, EdgeT)]
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> [(Int, b)]
G.lsuc Gr NodeT EdgeT
g Int
wacc_id
            TransNode out tr inp <- maybeToList $ G.lab g tn_id
            guard $ realConsumers tn_id == [wacc_id]
            pure (tn_id, out, tr, inp)
          trans_ids :: [Int]
trans_ids = ((Int, VName, ArrayTransform, VName) -> Int)
-> [(Int, VName, ArrayTransform, VName)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (\(Int
a, VName
_, ArrayTransform
_, VName
_) -> Int
a) [(Int, VName, ArrayTransform, VName)]
trans_preds
          trans_out_nms :: Names
trans_out_nms = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ ((Int, VName, ArrayTransform, VName) -> VName)
-> [(Int, VName, ArrayTransform, VName)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (\(Int
_, VName
out, ArrayTransform
_, VName
_) -> VName
out) [(Int, VName, ArrayTransform, VName)]
trans_preds
          -- The unique Screma SoacNode that feeds all TransNodes and has no
          -- other consumers besides those TransNodes.
          soac_preds :: [Int]
soac_preds = do
            (tn_id, _, _, _) <- [(Int, VName, ArrayTransform, VName)]
trans_preds
            (sn_id, _) <- G.lsuc g tn_id
            guard $ sn_id `notElem` trans_ids
            guard $ all (`elem` trans_ids) (realConsumers sn_id)
            pure sn_id
          prod_ids :: [Int]
prod_ids = [Int] -> [Int]
forall a. Ord a => [a] -> [a]
nubOrd [Int]
soac_preds
      case ([(Int, VName, ArrayTransform, VName)]
trans_preds, [Int]
prod_ids) of
        ((Int, VName, ArrayTransform, VName)
_ : [(Int, VName, ArrayTransform, VName)]
_, [Int
prod_id])
          | Just (SoacNode ArrayTransforms
ots1 Pat Type
pat1 soac :: SOAC SOACS
soac@(H.Screma {}) StmAux (ExpDec SOACS)
aux1) <- Gr NodeT EdgeT -> Int -> Maybe NodeT
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Maybe a
G.lab Gr NodeT EdgeT
g Int
prod_id,
            ArrayTransforms
ots1 ArrayTransforms -> ArrayTransforms -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayTransforms
forall a. Monoid a => a
mempty,
            (Input -> Bool) -> [Input] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((VName -> Names -> Bool
`notNameIn` Names
wacc_cons_nms) (VName -> Bool) -> (Input -> VName) -> Input -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> VName
H.inputArray) (SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
H.inputs SOAC SOACS
soac),
            Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Names -> Names -> Bool
namesIntersect Names
trans_out_nms Names
wacc_cons_nms ->
              [Int]
-> [(Int, VName, ArrayTransform, VName)]
-> Int
-> Pat Type
-> StmAux ()
-> SOAC SOACS
-> Pat Type
-> StmAux ()
-> [WithAccInput SOACS]
-> Lambda SOACS
-> m DepGraph
forall {a} {t}.
[Int]
-> [(a, VName, ArrayTransform, VName)]
-> Int
-> Pat Type
-> StmAux ()
-> SOAC SOACS
-> Pat Type
-> StmAux ()
-> [WithAccInput SOACS]
-> GLambda SOACS t
-> m DepGraph
attempt [Int]
trans_ids [(Int, VName, ArrayTransform, VName)]
trans_preds Int
prod_id Pat Type
pat1 StmAux ()
StmAux (ExpDec SOACS)
aux1 SOAC SOACS
soac Pat Type
Pat (LetDec SOACS)
pat2 StmAux ()
StmAux (ExpDec SOACS)
aux2 [WithAccInput SOACS]
w_inps Lambda SOACS
lam0
        ([(Int, VName, ArrayTransform, VName)], [Int])
_ -> DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  | Bool
otherwise = DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg
  where
    attempt :: [Int]
-> [(a, VName, ArrayTransform, VName)]
-> Int
-> Pat Type
-> StmAux ()
-> SOAC SOACS
-> Pat Type
-> StmAux ()
-> [WithAccInput SOACS]
-> GLambda SOACS t
-> m DepGraph
attempt [Int]
trans_ids [(a, VName, ArrayTransform, VName)]
trans_preds Int
prod_id Pat Type
pat1 StmAux ()
aux1 SOAC SOACS
soac Pat Type
pat2 StmAux ()
aux2 [WithAccInput SOACS]
w_inps GLambda SOACS t
lam0 = do
      let trans_info :: [(VName, ArrayTransform, VName)]
trans_info = ((a, VName, ArrayTransform, VName)
 -> (VName, ArrayTransform, VName))
-> [(a, VName, ArrayTransform, VName)]
-> [(VName, ArrayTransform, VName)]
forall a b. (a -> b) -> [a] -> [b]
map (\(a
_, VName
out, ArrayTransform
tr, VName
inp) -> (VName
out, ArrayTransform
tr, VName
inp)) [(a, VName, ArrayTransform, VName)]
trans_preds
      lam' <- Lambda SOACS -> m (Lambda SOACS)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda (Lambda SOACS -> m (Lambda SOACS))
-> (Builder SOACS [SubExpRes] -> m (Lambda SOACS))
-> Builder SOACS [SubExpRes]
-> m (Lambda SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< [LParam SOACS] -> Builder SOACS [SubExpRes] -> m (Lambda SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
[LParam rep] -> Builder rep [SubExpRes] -> m (Lambda rep)
runLambdaBuilder (GLambda SOACS t -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda SOACS t
lam0) (Builder SOACS [SubExpRes] -> m (Lambda SOACS))
-> Builder SOACS [SubExpRes] -> m (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ do
        soac' <- SOAC (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT
     SOACS
     (State VNameSource)
     (Exp (Rep (BuilderT SOACS (State VNameSource))))
forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m)) =>
SOAC (Rep m) -> m (Exp (Rep m))
H.toExp SOAC (Rep (BuilderT SOACS (State VNameSource)))
SOAC SOACS
soac
        addStm $ Let pat1 aux1 soac'
        forM_ trans_info $ \(VName
out, ArrayTransform
tr, VName
inp) -> do
          (tr_aux, tr_exp) <- ArrayTransform
-> VName
-> BuilderT SOACS (State VNameSource) (StmAux (), Exp SOACS)
forall (m :: * -> *) rep.
(Monad m, HasScope rep m) =>
ArrayTransform -> VName -> m (StmAux (), Exp rep)
H.transformToExp ArrayTransform
tr VName
inp
          auxing tr_aux $ letBindNames [out] tr_exp
        bodyBind $ lambdaBody lam0
      -- Run inner fusion. We always proceed with onSuccess because embedding
      -- the SoacNode + TransNodes into the WithAcc is itself a valid fusion
      -- step and avoids potential infinite loops from simplifyLambda hoisting
      -- a reshape back out.
      lam'' <- fst <$> doFusionInLambda lam'
      onSuccess trans_ids prod_id pat2 aux2 w_inps lam''

    onSuccess :: [Int]
-> Int
-> Pat Type
-> StmAux ()
-> [WithAccInput SOACS]
-> Lambda SOACS
-> m DepGraph
onSuccess [Int]
trans_ids Int
prod_id Pat Type
pat2 StmAux ()
aux2 [WithAccInput SOACS]
w_inps Lambda SOACS
lam'' = do
      m (Maybe NodeT) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Maybe NodeT) -> m ()) -> m (Maybe NodeT) -> m ()
forall a b. (a -> b) -> a -> b
$ NodeT -> m (Maybe NodeT)
fusedSomething (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 Pat Type
Pat (LetDec SOACS)
pat2 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]
w_inps Lambda SOACS
lam'')
      -- Rebuild the graph: remove absorbed nodes and rewire wacc's edges.
      -- G.context returns (in_adj, n, label, out_adj) where both adjacency
      -- lists are [(EdgeT, Node)]. G.lsuc/lpre are (Node, EdgeT).
      let to_remove :: [Int]
to_remove = Int
prod_id Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
trans_ids
          new_wacc :: NodeT
new_wacc = 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 Pat Type
Pat (LetDec SOACS)
pat2 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]
w_inps Lambda SOACS
lam''
          g' :: Gr NodeT EdgeT
g' = (Int -> Gr NodeT EdgeT -> Gr NodeT EdgeT)
-> Gr NodeT EdgeT -> [Int] -> Gr NodeT EdgeT
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Int -> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> gr a b
G.delNode Gr NodeT EdgeT
g [Int]
to_remove
          (Adj EdgeT
wacc_preds, Int
_, NodeT
_, Adj EdgeT
wacc_succs) = Gr NodeT EdgeT -> Int -> (Adj EdgeT, Int, NodeT, Adj EdgeT)
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Context a b
G.context Gr NodeT EdgeT
g Int
wacc_id
          -- Inherit producers of the absorbed nodes that are still in g'.
          removed_succs :: [(Int, EdgeT)]
removed_succs =
            [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a. Ord a => [a] -> [a]
nubOrd ([(Int, EdgeT)] -> [(Int, EdgeT)])
-> [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a b. (a -> b) -> a -> b
$
              (Int -> [(Int, EdgeT)]) -> [Int] -> [(Int, EdgeT)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (((Int, EdgeT) -> Bool) -> [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
`G.gelem` Gr NodeT EdgeT
g') (Int -> Bool) -> ((Int, EdgeT) -> Int) -> (Int, EdgeT) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, EdgeT) -> Int
forall a b. (a, b) -> a
fst) ([(Int, EdgeT)] -> [(Int, EdgeT)])
-> (Int -> [(Int, EdgeT)]) -> Int -> [(Int, EdgeT)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gr NodeT EdgeT -> Int -> [(Int, EdgeT)]
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> [(Int, b)]
G.lsuc Gr NodeT EdgeT
g) [Int]
to_remove
          new_succs :: [(Int, EdgeT)]
new_succs =
            [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a. Ord a => [a] -> [a]
nubOrd ([(Int, EdgeT)] -> [(Int, EdgeT)])
-> [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a b. (a -> b) -> a -> b
$
              [(Int, EdgeT)]
removed_succs
                [(Int, EdgeT)] -> [(Int, EdgeT)] -> [(Int, EdgeT)]
forall a. [a] -> [a] -> [a]
++ ((EdgeT, Int) -> (Int, EdgeT)) -> Adj EdgeT -> [(Int, EdgeT)]
forall a b. (a -> b) -> [a] -> [b]
map (\(EdgeT
e, Int
n) -> (Int
n, EdgeT
e)) (((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
`G.gelem` Gr NodeT EdgeT
g') (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
wacc_succs)
          new_preds :: Adj EdgeT
new_preds = ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Gr NodeT EdgeT -> Bool
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> Bool
`G.gelem` Gr NodeT EdgeT
g') (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
wacc_preds
          g'' :: Gr NodeT EdgeT
g'' = LNode NodeT -> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall (gr :: * -> * -> *) a b.
DynGraph gr =>
LNode a -> gr a b -> gr a b
G.insNode (Int
wacc_id, NodeT
new_wacc) (Gr NodeT EdgeT -> Gr NodeT EdgeT)
-> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall a b. (a -> b) -> a -> b
$ Int -> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> gr a b
G.delNode Int
wacc_id Gr NodeT EdgeT
g'
          g''' :: Gr NodeT EdgeT
g''' = ((EdgeT, Int) -> Gr NodeT EdgeT -> Gr NodeT EdgeT)
-> Gr NodeT EdgeT -> Adj EdgeT -> Gr NodeT EdgeT
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(EdgeT
e, Int
n) Gr NodeT EdgeT
gr -> LEdge EdgeT -> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall (gr :: * -> * -> *) b a.
DynGraph gr =>
LEdge b -> gr a b -> gr a b
G.insEdge (Int
n, Int
wacc_id, EdgeT
e) Gr NodeT EdgeT
gr) Gr NodeT EdgeT
g'' Adj EdgeT
new_preds
          g'''' :: Gr NodeT EdgeT
g'''' = ((Int, EdgeT) -> Gr NodeT EdgeT -> Gr NodeT EdgeT)
-> Gr NodeT EdgeT -> [(Int, EdgeT)] -> Gr NodeT EdgeT
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Int
n, EdgeT
e) Gr NodeT EdgeT
gr -> LEdge EdgeT -> Gr NodeT EdgeT -> Gr NodeT EdgeT
forall (gr :: * -> * -> *) b a.
DynGraph gr =>
LEdge b -> gr a b -> gr a b
G.insEdge (Int
wacc_id, Int
n, EdgeT
e) Gr NodeT EdgeT
gr) Gr NodeT EdgeT
g''' [(Int, EdgeT)]
new_succs
      DepGraph -> m DepGraph
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DepGraph
dg {dgGraph = g''''}

-------------------------------
--- simple helper functions ---
-------------------------------

substInSEs :: M.Map VName VName -> [SubExp] -> [SubExp]
substInSEs :: Map VName VName -> [SubExp] -> [SubExp]
substInSEs Map VName VName
vtab = (SubExp -> SubExp) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> SubExp
substInSE
  where
    substInSE :: SubExp -> SubExp
substInSE (Var VName
x)
      | Just VName
y <- VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Map VName VName
vtab = VName -> SubExp
Var VName
y
    substInSE SubExp
z = SubExp
z

equivLambda ::
  M.Map VName VName ->
  Lambda SOACS ->
  Lambda SOACS ->
  Bool
equivLambda :: Map VName VName -> Lambda SOACS -> Lambda SOACS -> Bool
equivLambda Map VName VName
stab Lambda SOACS
lam1 Lambda SOACS
lam2
  | ([Param Type]
ps1, [Param Type]
ps2) <- (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam1, Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam2),
    ([VName]
nms1, [VName]
nms2) <- ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
ps1, (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
ps2),
    (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Param dec -> dec
paramDec [Param Type]
ps1 [Type] -> [Type] -> Bool
forall a. Eq a => a -> a -> Bool
== (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Param dec -> dec
paramDec [Param Type]
ps2,
    (Param Type -> Attrs) -> [Param Type] -> [Attrs]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Attrs
forall dec. Param dec -> Attrs
paramAttrs [Param Type]
ps1 [Attrs] -> [Attrs] -> Bool
forall a. Eq a => a -> a -> Bool
== (Param Type -> Attrs) -> [Param Type] -> [Attrs]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Attrs
forall dec. Param dec -> Attrs
paramAttrs [Param Type]
ps2,
    Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam1 [Type] -> [Type] -> Bool
forall a. Eq a => a -> a -> Bool
== Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam2,
    (GBody SOACS SubExpRes
bdy1, GBody SOACS SubExpRes
bdy2) <- (Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam1, Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam2),
    GBody SOACS SubExpRes -> BodyDec SOACS
forall rep res. GBody rep res -> BodyDec rep
bodyDec GBody SOACS SubExpRes
bdy1 () -> () -> Bool
forall a. Eq a => a -> a -> Bool
== GBody SOACS SubExpRes -> BodyDec SOACS
forall rep res. GBody rep res -> BodyDec rep
bodyDec GBody SOACS SubExpRes
bdy2 =
      let insert :: Map k a -> (a, k) -> Map k a
insert Map k a
tab (a
x, k
k) = k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k a
x Map k a
tab
          stab' :: Map VName VName
stab' = (Map VName VName -> (VName, VName) -> Map VName VName)
-> Map VName VName -> [(VName, VName)] -> Map VName VName
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map VName VName -> (VName, VName) -> Map VName VName
forall {k} {a}. Ord k => Map k a -> (a, k) -> Map k a
insert Map VName VName
stab ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$ [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
nms1 [VName]
nms2
          fStm :: (Map VName VName, Bool)
-> (Stm SOACS, Stm SOACS) -> (Map VName VName, Bool)
fStm (Map VName VName
vtab, Bool
False) (Stm SOACS, Stm SOACS)
_ = (Map VName VName
vtab, Bool
False)
          fStm (Map VName VName
vtab, Bool
True) (Stm SOACS
s1, Stm SOACS
s2) = Map VName VName
-> Stm SOACS -> Stm SOACS -> (Map VName VName, Bool)
equivStm Map VName VName
vtab Stm SOACS
s1 Stm SOACS
s2
          (Map VName VName
stab'', Bool
success) =
            ((Map VName VName, Bool)
 -> (Stm SOACS, Stm SOACS) -> (Map VName VName, Bool))
-> (Map VName VName, Bool)
-> [(Stm SOACS, Stm SOACS)]
-> (Map VName VName, Bool)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Map VName VName, Bool)
-> (Stm SOACS, Stm SOACS) -> (Map VName VName, Bool)
fStm (Map VName VName
stab', Bool
True) ([(Stm SOACS, Stm SOACS)] -> (Map VName VName, Bool))
-> [(Stm SOACS, Stm SOACS)] -> (Map VName VName, Bool)
forall a b. (a -> b) -> a -> b
$
              [Stm SOACS] -> [Stm SOACS] -> [(Stm SOACS, Stm SOACS)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList (GBody SOACS SubExpRes -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms GBody SOACS SubExpRes
bdy1)) ([Stm SOACS] -> [(Stm SOACS, Stm SOACS)])
-> [Stm SOACS] -> [(Stm SOACS, Stm SOACS)]
forall a b. (a -> b) -> a -> b
$
                Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList (GBody SOACS SubExpRes -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms GBody SOACS SubExpRes
bdy2)
          sres2 :: [SubExp]
sres2 = Map VName VName -> [SubExp] -> [SubExp]
substInSEs Map VName VName
stab'' ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp ([SubExpRes] -> [SubExp]) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ GBody SOACS SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult GBody SOACS SubExpRes
bdy2
       in Bool
success Bool -> Bool -> Bool
&& (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp (GBody SOACS SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult GBody SOACS SubExpRes
bdy1) [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== [SubExp]
sres2
equivLambda Map VName VName
_ Lambda SOACS
_ Lambda SOACS
_ =
  Bool
False

equivStm ::
  M.Map VName VName ->
  Stm SOACS ->
  Stm SOACS ->
  (M.Map VName VName, Bool)
equivStm :: Map VName VName
-> Stm SOACS -> Stm SOACS -> (Map VName VName, Bool)
equivStm
  Map VName VName
stab
  (Let Pat (LetDec SOACS)
pat1 StmAux (ExpDec SOACS)
aux1 (BasicOp (BinOp BinOp
bop1 SubExp
se11 SubExp
se12)))
  (Let Pat (LetDec SOACS)
pat2 StmAux (ExpDec SOACS)
aux2 (BasicOp (BinOp BinOp
bop2 SubExp
se21 SubExp
se22)))
    | [SubExp
se11, SubExp
se12] [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== Map VName VName -> [SubExp] -> [SubExp]
substInSEs Map VName VName
stab [SubExp
se21, SubExp
se22],
      ([PatElem Type]
pels1, [PatElem Type]
pels2) <- (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat1, Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat2),
      (PatElem Type -> Type) -> [PatElem Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map PatElem Type -> Type
forall dec. PatElem dec -> dec
patElemDec [PatElem Type]
pels1 [Type] -> [Type] -> Bool
forall a. Eq a => a -> a -> Bool
== (PatElem Type -> Type) -> [PatElem Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map PatElem Type -> Type
forall dec. PatElem dec -> dec
patElemDec [PatElem Type]
pels2,
      BinOp
bop1 BinOp -> BinOp -> Bool
forall a. Eq a => a -> a -> Bool
== BinOp
bop2 Bool -> Bool -> Bool
&& StmAux ()
StmAux (ExpDec SOACS)
aux1 StmAux () -> StmAux () -> Bool
forall a. Eq a => a -> a -> Bool
== StmAux ()
StmAux (ExpDec SOACS)
aux2 =
        let stab_new :: Map VName VName
stab_new =
              [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$
                [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((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]
pels2) ((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]
pels1)
         in (Map VName VName -> Map VName VName -> Map VName VName
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Map VName VName
stab_new Map VName VName
stab, Bool
True)
-- To Be Continued ...
equivStm Map VName VName
vtab Stm SOACS
_ Stm SOACS
_ = (Map VName VName
vtab, Bool
False)