futhark-0.25.29: src/Futhark/Optimise/Fusion/RulesWithAccs.hs
{-# LANGUAGE Strict #-}
-- | This module consists of rules for fusion
-- that involves WithAcc constructs.
-- Currently, we support two non-trivial
-- transformations:
-- I. map-flatten-scatter: a map nest produces
-- multi-dimensional index and values arrays
-- that are then flattened and used in a
-- scatter consumer. Such pattern can be fused
-- by re-writing the scatter by means of a WithAcc
-- containing a map-nest, thus eliminating the flatten
-- operations. The obtained WithAcc can then be fused
-- with the producer map nest, e.g., benefiting intra-group
-- kernels. The eloquent target for this rule is
-- an efficient implementation of radix-sort.
--
-- 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
( ruleMFScat,
tryFuseWithAccs,
)
where
import Control.Monad
import Data.Graph.Inductive.Graph qualified as G
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.HORep.SOAC qualified as H
import Futhark.Construct
import Futhark.IR.SOACS hiding (SOAC (..))
import Futhark.IR.SOACS qualified as F
import Futhark.Optimise.Fusion.GraphRep
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
se0 :: SubExp
se0 = intConst Int64 0
se1 :: SubExp
se1 = intConst Int64 1
-------------------------------------
--- I. Map-Flatten-Scatter Fusion ---
-------------------------------------
-- helper data structures
type IotaInp = ((VName, LParam SOACS), (SubExp, SubExp, SubExp, IntType))
-- ^ ((array-name, lambda param), (len, start, stride, Int64))
type RshpInp = ((VName, LParam SOACS), (Shape, Shape, Type))
-- ^ ((array-name, lambda param), (flat-shape, unflat-shape, elem-type))
-- | Implements a specialized rule for fusing a pattern
-- formed by a map o flatten o scatter, i.e.,
-- let (inds, vals) = map-nest f inps
-- (finds, fvals) = (flatten inds, flatten vals)
-- let res = scatter res0 finds fvals
-- where inds & vals have higher rank than finds & fvals.
ruleMFScat ::
(HasScope SOACS m, MonadFreshNames m) =>
DepNode ->
DepGraph ->
m (Maybe DepGraph)
ruleMFScat node_to_fuse dg@DepGraph {dgGraph = g}
| soac_nodeT <- snd node_to_fuse,
scat_node_id <- nodeFromLNode node_to_fuse,
SoacNode node_out_trsfs scat_pat scat_soac scat_aux <- soac_nodeT,
H.nullTransforms node_out_trsfs,
-- \^ for simplicity we do not allow transforms on scatter's result.
H.Scatter _len scat_inp scat_out scat_lam <- scat_soac,
-- \^ get the scatter
scat_trsfs <- map H.inputTransforms (H.inputs scat_soac),
-- \^ get the transforms on the input
any (/= mempty) scat_trsfs,
scat_ctx <- G.context g scat_node_id,
(out_deps, _, _, inp_deps) <- scat_ctx,
cons_deps <- filter (isCons . fst) inp_deps,
drct_deps <- filter (isDep . fst) inp_deps,
cons_ctxs <- map (G.context g . snd) cons_deps,
drct_ctxs <- map (G.context g . snd) drct_deps,
_cons_nTs <- map getNodeTfromCtx cons_ctxs, -- not used!!
drct_tups0 <- mapMaybe (pairUp (zip drct_ctxs (map fst drct_deps))) scat_inp,
length drct_tups0 == length scat_inp,
-- \^ checks that all direct dependencies are also array
-- inputs to scatter
(t1s, t2s) <- unzip drct_tups0,
drct_tups <- zip t1s $ zip t2s (lambdaParams scat_lam),
(ctxs_iots, drct_iots) <- unzip $ filter (isIota . snd . fst . snd) drct_tups,
(ctxs_rshp, drct_rshp) <- unzip $ filter (not . isIota . snd . fst . snd) drct_tups,
length drct_iots + length drct_rshp == length scat_inp,
-- \^ direct dependencies are either flatten reshapes or iotas.
rep_iotas <- mapMaybe getRepIota drct_iots,
length rep_iotas == length drct_iots,
rep_rshps_certs <- mapMaybe getRepRshpArr drct_rshp,
(rep_rshps, certs_rshps) <- unzip rep_rshps_certs,
-- \^ gather the representations for the iotas and reshapes, that use
-- the helper types `IotaInp` and `RshpInp`
not (null rep_rshps),
-- \^ at least one flatten-reshaped array
length rep_rshps == length drct_rshp,
(_, (s1, s2, _)) : _ <- rep_rshps,
all (\(_, (s1', s2', _)) -> s1 == s1' && s2 == s2') rep_rshps,
-- \^ Check that all unflatten shape dimensions are the same,
-- so that we can construct a map nest;
-- check profitability, which is conservatively defined as all
-- the reshaped and consumer arrays are used solely by the
-- scatter AND all reshape dependencies originate in the same
-- map.
checkSafeAndProfitable dg scat_node_id ctxs_rshp cons_ctxs = do
-- generate the withAcc statement
let cons_patels_outs = zip (patElems scat_pat) scat_out
wacc_stm <- mkWithAccStm rep_iotas rep_rshps cons_patels_outs scat_aux scat_lam
let all_cert_rshp = mconcat certs_rshps
aux = stmAux wacc_stm
aux' = aux {stmAuxCerts = all_cert_rshp <> stmAuxCerts aux}
wacc_stm' = wacc_stm {stmAux = aux'}
-- get the input deps of iotas
fiot acc (_, _, _, inp_deps_iot) =
acc <> inp_deps_iot
deps_of_iotas = foldl fiot mempty ctxs_iots
--
iota_nms = namesFromList $ map (fst . fst) rep_iotas
inp_deps_wo_iotas = filter ((`notNameIn` iota_nms) . getName . fst) inp_deps
-- generate a new node for the with-acc-stmt and its associated context:
-- add the inp-deps of iotas but remove the iota themselves from deps.
new_withacc_nT = StmNode wacc_stm'
inp_deps' = inp_deps_wo_iotas <> deps_of_iotas
new_withacc_ctx = (out_deps, scat_node_id, new_withacc_nT, inp_deps')
-- construct the new WithAcc node/graph; do we need to use `fusedSomething` ??
new_node = G.node' new_withacc_ctx
dg' = dg {dgGraph = new_withacc_ctx G.& G.delNodes [new_node] g}
-- result
pure $ Just dg'
where
--
getNodeTfromCtx (_, _, nT, _) = nT
findCtxOf ctxes nm
| [ctxe] <- filter (\x -> nm == getName (snd x)) ctxes =
Just ctxe
findCtxOf _ _ = Nothing
pairUp :: [(DepContext, EdgeT)] -> H.Input -> Maybe (DepContext, (H.Input, NodeT))
pairUp ctxes inp@(H.Input _arrtrsfs nm _tp)
| Just (ctx@(_, _, nT, _), _) <- findCtxOf ctxes nm =
Just (ctx, (inp, nT))
pairUp _ _ = Nothing
--
isIota :: NodeT -> Bool
isIota (StmNode (Let _ _ (BasicOp (Iota {})))) = True
isIota _ = False
--
getRepIota :: ((H.Input, NodeT), LParam SOACS) -> Maybe IotaInp
getRepIota ((H.Input iottrsf arr_nm _arr_tp, nt), farg)
| mempty == iottrsf,
StmNode (Let _ _ (BasicOp (Iota n x s Int64))) <- nt =
Just ((arr_nm, farg), (n, x, s, Int64))
getRepIota _ = Nothing
--
getRepRshpArr :: ((H.Input, NodeT), LParam SOACS) -> Maybe (RshpInp, Certs)
getRepRshpArr ((H.Input outtrsf arr_nm arr_tp, _nt), farg)
| rshp_trsfm H.:< other_trsfms <- H.viewf outtrsf,
(H.Reshape c ReshapeArbitrary shp_flat) <- rshp_trsfm,
other_trsfms == mempty,
eltp <- paramDec farg,
Just shp_flat' <- checkShp eltp shp_flat,
Array _ptp shp_unflat _ <- arr_tp,
Just shp_unflat' <- checkShp eltp shp_unflat,
shapeRank shp_flat' == 1,
shapeRank shp_flat' < shapeRank shp_unflat' =
Just (((arr_nm, farg), (shp_flat', shp_unflat', eltp)), c)
getRepRshpArr _ = Nothing
--
checkShp (Prim _) shp_arr = Just shp_arr
checkShp (Array _ptp shp_elm _) shp_arr =
let dims_elm = shapeDims shp_elm
dims_arr = shapeDims shp_arr
(m, n) = (length dims_elm, length dims_arr)
shp' = Shape $ take (n - m) dims_arr
dims_com = drop (n - m) dims_arr
in if all (uncurry (==)) (zip dims_com dims_elm)
then Just shp'
else Nothing
checkShp _ _ = Nothing
-- default fails:
ruleMFScat _ _ = pure Nothing
checkSafeAndProfitable :: DepGraph -> G.Node -> [DepContext] -> [DepContext] -> Bool
checkSafeAndProfitable dg scat_node_id ctxs_rshp@(_ : _) ctxs_cons =
let all_deps = concatMap (\(x, _, _, _) -> x) $ ctxs_rshp ++ ctxs_cons
prof1 = all (\(_, dep_id) -> dep_id == scat_node_id) all_deps
-- \^ scatter is the sole target to all consume & unflatten-reshape deps
(_, map_node_id, map_nT, _) = head ctxs_rshp
prof2 = all (\(_, nid, _, _) -> nid == map_node_id) ctxs_rshp
prof3 = isMap map_nT
-- \^ all reshapes come from the same node, which is a map
safe = vFusionFeasability dg map_node_id scat_node_id
in safe && prof1 && prof2 && prof3
where
isMap nT
| SoacNode out_trsfs _pat soac _ <- nT,
H.Screma _ _ form <- soac,
ScremaForm _ [] [] <- form =
H.nullTransforms out_trsfs
isMap _ = False
checkSafeAndProfitable _ _ _ _ = False
-- | produces the withAcc statement that constitutes the translation of
-- the scater o flatten o map composition in which the map inputs are
-- reshaped in the same way
mkWithAccStm ::
(HasScope SOACS m, MonadFreshNames m) =>
[IotaInp] ->
[RshpInp] ->
[(PatElem (LetDec SOACS), (Shape, Int, VName))] ->
StmAux (ExpDec SOACS) ->
Lambda SOACS ->
m (Stm SOACS)
mkWithAccStm iota_inps rshp_inps cons_patels_outs scatter_aux scatter_lam
-- iotas are assumed to operate on Int64 values
-- ToDo: maybe simplify rshp_inps
-- check that the unflat shape is the same across reshapes
-- check that the rank of the unflatten shape is higher than the flatten
| rshp_inp : _ <- rshp_inps,
(_, (_, s_unflat, _)) <- rshp_inp,
(_ : _) <- shapeDims s_unflat = do
--
(cert_params, acc_params) <- fmap unzip $
forM cons_patels_outs $ \(patel, (shp, _, nm)) -> do
cert_param <- newParam "acc_cert_p" $ Prim Unit
let arr_tp = patElemType patel
acc_tp = stripArray (shapeRank shp) arr_tp
acc_param <-
newParam (baseString nm) $
Acc (paramName cert_param) shp [acc_tp] NoUniqueness
pure (cert_param, acc_param)
let cons_params_outs = zip acc_params $ map snd cons_patels_outs
acc_bdy <- mkWithAccBdy s_unflat iota_inps rshp_inps cons_params_outs scatter_lam
let withacc_lam =
Lambda
{ lambdaParams = cert_params ++ acc_params,
lambdaReturnType = map paramDec acc_params,
lambdaBody = acc_bdy
}
withacc_inps = map (\(_, (shp, _, nm)) -> (shp, [nm], Nothing)) cons_patels_outs
withacc_pat = Pat $ map fst cons_patels_outs
stm =
Let withacc_pat scatter_aux $
WithAcc withacc_inps withacc_lam
pure stm
mkWithAccStm _ _ _ _ _ =
error "Unreachable case reached!"
-- | Wrapper function for constructing the body of the withAcc
-- translation of the scatter
mkWithAccBdy ::
(HasScope SOACS m, MonadFreshNames m) =>
Shape ->
[IotaInp] ->
[RshpInp] ->
[(LParam SOACS, (Shape, Int, VName))] ->
Lambda SOACS ->
m (Body SOACS)
mkWithAccBdy shp iota_inps rshp_inps cons_params_outs scat_lam = do
let cons_ps = map fst cons_params_outs
scat_res_info = map snd cons_params_outs
static_arg = (iota_inps, rshp_inps, scat_res_info, scat_lam)
mkParam ((nm, _), (_, s, t)) = Param mempty nm (arrayOfShape t s)
rshp_ps = map mkParam rshp_inps
mkWithAccBdy' static_arg (shapeDims shp) [] [] rshp_ps cons_ps
-- | builds a body that essentially consists of a map-nest with accumulators,
-- i.e., one level for each level of the unflatten shape of scatter's reshaped
-- input arrays
mkWithAccBdy' ::
(HasScope SOACS m, MonadFreshNames m) =>
([IotaInp], [RshpInp], [(Shape, Int, VName)], Lambda SOACS) ->
[SubExp] ->
[SubExp] ->
[VName] ->
[LParam SOACS] ->
[LParam SOACS] ->
m (Body SOACS)
-- | the base case below addapts the scatter's lambda
mkWithAccBdy' static_arg [] dims_rev iot_par_nms rshp_ps cons_ps = do
let (iota_inps, rshp_inps, scat_res_info, scat_lam) = static_arg
tp_int = Prim $ IntType Int64
scope <- askScope
runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do
-- handle iota args
let strides_rev = scanl (*) (pe64 se1) $ map pe64 dims_rev
strides = tail $ reverse strides_rev
prods = zipWith (*) (map le64 iot_par_nms) strides
i_pe = sum prods
i_norm <- letExp "iota_norm_arg" =<< toExp i_pe
forM_ iota_inps $ \arg -> do
let ((_, i_par), (_, b, s, _)) = arg
i_new <- letExp "tmp" =<< toExp (pe64 b + le64 i_norm * pe64 s)
letBind (Pat [PatElem (paramName i_par) tp_int]) $ BasicOp $ SubExp $ Var i_new
-- handle rshp args
let rshp_lam_args = map (snd . fst) rshp_inps
forM_ (zip rshp_lam_args rshp_ps) $ \(old_par, new_par) -> do
let pat = Pat [PatElem (paramName old_par) (paramDec old_par)]
letBind pat $ BasicOp $ SubExp $ Var $ paramName new_par
-- add the body of the scatter's lambda
mapM_ addStm $ bodyStms $ lambdaBody scat_lam
-- add the withAcc update statements
let iv_ses = groupScatterResults' scat_res_info $ bodyResult $ lambdaBody scat_lam
res_nms <-
forM (zip cons_ps iv_ses) $ \(cons_p, (i_ses, v_se)) -> do
-- i_ses is a list
let f nm_in i_se =
letExp (baseString nm_in) $ BasicOp $ UpdateAcc Safe nm_in [resSubExp i_se] [resSubExp v_se]
foldM f (paramName cons_p) i_ses
let lam_certs = foldMap resCerts $ bodyResult $ lambdaBody scat_lam
pure $ map (SubExpRes lam_certs . Var) res_nms
-- \| the recursive case builds a call to a map soac.
mkWithAccBdy' static_arg (dim : dims) dims_rev iot_par_nms rshp_ps cons_ps = do
scope <- askScope
runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do
iota_arr <- letExp "iota_arr" $ BasicOp $ Iota dim se0 se1 Int64
iota_p <- newParam "iota_arg" $ Prim $ IntType Int64
rshp_ps' <- forM (zip [0 .. length rshp_ps - 1] (map paramDec rshp_ps)) $
\(i, arr_tp) ->
newParam ("rshp_arg_" ++ show i) $ stripArray 1 arr_tp
cons_ps' <- forM (zip [0 .. length cons_ps - 1] (map paramDec cons_ps)) $
\(i, arr_tp) ->
newParam ("acc_arg_" ++ show i) arr_tp
map_lam_bdy <-
mkWithAccBdy' static_arg dims (dim : dims_rev) (iot_par_nms ++ [paramName iota_p]) rshp_ps' cons_ps'
let map_lam = Lambda (rshp_ps' ++ [iota_p] ++ cons_ps') (map paramDec cons_ps') map_lam_bdy
map_inps = map paramName rshp_ps ++ [iota_arr] ++ map paramName cons_ps
map_soac = F.Screma dim map_inps $ ScremaForm map_lam [] []
res_nms <- letTupExp "acc_res" $ Op map_soac
pure $ map (subExpRes . Var) res_nms
---------------------------------------------------
--- 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
type AccTup =
( [PatElem (LetDec SOACS)],
WithAccInput SOACS,
LParam SOACS,
LParam SOACS,
(VName, Certs)
)
accTup1 :: AccTup -> [PatElem (LetDec SOACS)]
accTup1 (a, _, _, _, _) = a
accTup2 :: AccTup -> WithAccInput SOACS
accTup2 (_, a, _, _, _) = a
accTup3 :: AccTup -> LParam SOACS
accTup3 (_, _, a, _, _) = a
accTup4 :: AccTup -> LParam SOACS
accTup4 (_, _, _, a, _) = a
accTup5 :: AccTup -> (VName, Certs)
accTup5 (_, _, _, _, a) = a
-- | 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
infusible
(Let pat1 aux1 (WithAcc w_inps1 lam1))
(Let pat2 aux2 (WithAcc w_inps2 lam2))
| (pat1_els, pat2_els) <- (patElems pat1, patElems pat2),
(acc_tup1, other_pr1) <- groupAccs pat1_els w_inps1 lam1,
(acc_tup2, other_pr2) <- groupAccs pat2_els w_inps2 lam2,
(tup_common, acc_tup1', acc_tup2') <-
groupCommonAccs acc_tup1 acc_tup2,
-- safety 0: make sure that the accs from acc_tup1' and
-- acc_tup2' do not overlap
pnms_1' <- map patElemName $ concatMap (\(nms, _, _, _, _) -> nms) acc_tup1',
winp_2' <- concatMap (\(_, (_, nms, _), _, _, _) -> nms) acc_tup2',
not $ namesIntersect (namesFromList pnms_1') (namesFromList winp_2'),
-- safety 1: we have already determined the commons;
-- now we also need to check NOT-IN FV(lam2)
not $ namesIntersect (namesFromList pnms_1') (freeIn lam2),
-- safety 2:
-- bs <- map patElemName $ concatMap accTup1 acc_tup1,
bs <- map patElemName $ concatMap (accTup1 . fst) tup_common,
all (`notElem` infusible) bs,
-- safety 3:
cs <- namesFromList $ concatMap ((\(_, xs, _) -> xs) . accTup2) acc_tup2,
all ((`notNameIn` cs) . patElemName . fst) other_pr1 = Just $ do
let getCertPairs (t1, t2) = (paramName (accTup3 t2), paramName (accTup3 t1))
tab_certs = M.fromList $ map getCertPairs tup_common
lam2_bdy' = substituteNames tab_certs (lambdaBody lam2)
rcrt_params = map (accTup3 . fst) tup_common ++ map accTup3 acc_tup1' ++ map accTup3 acc_tup2'
racc_params = map (accTup4 . fst) tup_common ++ map accTup4 acc_tup1' ++ map accTup4 acc_tup2'
(comm_res_nms, comm_res_certs2) = unzip $ map (accTup5 . snd) tup_common
(_, comm_res_certs1) = unzip $ map (accTup5 . fst) tup_common
com_res_certs = zipWith (\x y -> Certs (unCerts x ++ unCerts y)) comm_res_certs1 comm_res_certs2
bdyres_certs = com_res_certs ++ map (snd . accTup5) (acc_tup1' ++ acc_tup2')
bdyres_accse = map Var comm_res_nms ++ map (Var . fst . accTup5) (acc_tup1' ++ acc_tup2')
bdy_res_accs = zipWith SubExpRes bdyres_certs bdyres_accse
bdy_res_others = map snd $ other_pr1 ++ other_pr2
scope <- askScope
lam_bdy <-
runBodyBuilder $ do
localScope (scope <> scopeOfLParams (rcrt_params ++ racc_params)) $ do
-- add the stms of lam1
mapM_ addStm $ stmsToList $ bodyStms $ lambdaBody lam1
-- add the copy stms for the common accumulator
forM_ tup_common $ \(tup1, tup2) -> do
let (lpar1, lpar2) = (accTup4 tup1, accTup4 tup2)
((nm1, _), nm2, tp_acc) = (accTup5 tup1, paramName lpar2, paramDec lpar1)
letBind (Pat [PatElem nm2 tp_acc]) $ BasicOp $ SubExp $ Var nm1
-- add copy stms to bring in scope x1 ... xq
forM_ other_pr1 $ \(pat_elm, bdy_res) -> do
let (nm, se, tp) = (patElemName pat_elm, resSubExp bdy_res, patElemType pat_elm)
certifying (resCerts bdy_res) $
letBind (Pat [PatElem nm tp]) $
BasicOp (SubExp se)
-- add the statements of lam2 (in which the acc-certificates have been substituted)
mapM_ addStm $ stmsToList $ bodyStms lam2_bdy'
-- build the result of body
pure $ bdy_res_accs ++ bdy_res_others
let tp_res_other = map (patElemType . fst) (other_pr1 ++ other_pr2)
res_lam =
Lambda
{ lambdaParams = rcrt_params ++ racc_params,
lambdaBody = lam_bdy,
lambdaReturnType = map paramDec racc_params ++ tp_res_other
}
res_lam' <- renameLambda res_lam
let res_pat =
concatMap (accTup1 . snd) tup_common
++ concatMap accTup1 (acc_tup1' ++ acc_tup2')
++ map fst (other_pr1 ++ other_pr2)
res_w_inps = map (accTup2 . fst) tup_common ++ map accTup2 (acc_tup1' ++ acc_tup2')
res_w_inps' <- mapM renameLamInWAccInp res_w_inps
pure $ Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam'
where
-- local helpers:
groupAccs ::
[PatElem (LetDec SOACS)] ->
[WithAccInput SOACS] ->
Lambda SOACS ->
([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccs pat_els wacc_inps wlam =
let lam_params = lambdaParams wlam
n = length lam_params
(lam_par_crts, lam_par_accs) = splitAt (n `div` 2) lam_params
lab_res_ses = bodyResult $ lambdaBody wlam
in groupAccsHlp pat_els wacc_inps lam_par_crts lam_par_accs lab_res_ses
groupAccsHlp ::
[PatElem (LetDec SOACS)] ->
[WithAccInput SOACS] ->
[LParam SOACS] ->
[LParam SOACS] ->
[SubExpRes] ->
([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
groupAccsHlp pat_els [] [] [] lam_res_ses
| length pat_els == length lam_res_ses =
([], zip pat_els lam_res_ses)
groupAccsHlp
pat_els
(winp@(_, inp, _) : wacc_inps)
(par_crt : lam_par_crts)
(par_acc : lam_par_accs)
(res_se : lam_res_ses)
| n <- length inp,
(n <= length pat_els) && (n <= (1 + length lam_res_ses)),
Var res_nm <- resSubExp res_se =
let (pat_els_cur, pat_els') = splitAt n pat_els
(rec1, rec2) = groupAccsHlp pat_els' wacc_inps lam_par_crts lam_par_accs lam_res_ses
in ((pat_els_cur, winp, par_crt, par_acc, (res_nm, resCerts res_se)) : rec1, rec2)
groupAccsHlp _ _ _ _ _ =
error "Unreachable case reached in groupAccsHlp!"
--
groupCommonAccs :: [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
groupCommonAccs [] tup_accs2 =
([], [], tup_accs2)
groupCommonAccs (tup_acc1 : tup_accs1) tup_accs2
| commons2 <- filter (matchingAccTup tup_acc1) tup_accs2,
length commons2 <= 1 =
let (rec1, rec2, rec3) =
groupCommonAccs tup_accs1 $
if null commons2
then tup_accs2
else filter (not . matchingAccTup tup_acc1) tup_accs2
in if null commons2
then (rec1, tup_acc1 : rec2, rec3)
else ((tup_acc1, head commons2) : rec1, tup_accs1, rec3)
groupCommonAccs _ _ =
error "Unreachable case reached in groupCommonAccs!"
renameLamInWAccInp (shp, inps, Just (lam, se)) = do
lam' <- renameLambda lam
pure (shp, inps, Just (lam', se))
renameLamInWAccInp winp = pure winp
--
tryFuseWithAccs _ _ _ =
Nothing
-------------------------------
--- simple helper functions ---
-------------------------------
equivLambda ::
M.Map VName VName ->
Lambda SOACS ->
Lambda SOACS ->
Bool
equivLambda stab lam1 lam2
| (ps1, ps2) <- (lambdaParams lam1, lambdaParams lam2),
(nms1, nms2) <- (map paramName ps1, map paramName ps2),
map paramDec ps1 == map paramDec ps2,
map paramAttrs ps1 == map paramAttrs ps2,
lambdaReturnType lam1 == lambdaReturnType lam2,
(bdy1, bdy2) <- (lambdaBody lam1, lambdaBody lam2),
bodyDec bdy1 == bodyDec bdy2 =
let insert tab (x, k) = M.insert k x tab
stab' = foldl insert stab $ zip nms1 nms2
fStm (vtab, False) _ = (vtab, False)
fStm (vtab, True) (s1, s2) = equivStm vtab s1 s2
(stab'', success) =
foldl fStm (stab', True) $
zip (stmsToList (bodyStms bdy1)) $
stmsToList (bodyStms bdy2)
sres2 = substInSEs stab'' $ map resSubExp $ bodyResult bdy2
in success && map resSubExp (bodyResult bdy1) == sres2
equivLambda _ _ _ =
False
equivStm ::
M.Map VName VName ->
Stm SOACS ->
Stm SOACS ->
(M.Map VName VName, Bool)
equivStm
stab
(Let pat1 aux1 (BasicOp (BinOp bop1 se11 se12)))
(Let pat2 aux2 (BasicOp (BinOp bop2 se21 se22)))
| [se11, se12] == substInSEs stab [se21, se22],
(pels1, pels2) <- (patElems pat1, patElems pat2),
map patElemDec pels1 == map patElemDec pels2,
bop1 == bop2 && aux1 == aux2 =
let stab_new =
M.fromList $
zip (map patElemName pels2) (map patElemName pels1)
in (M.union stab_new stab, True)
-- To Be Continued ...
equivStm vtab _ _ = (vtab, False)
matchingAccTup :: AccTup -> AccTup -> Bool
matchingAccTup
(pat_els1, (shp1, _winp_arrs1, mlam1), _, _, _)
(_, (shp2, winp_arrs2, mlam2), _, _, _) =
shapeDims shp1 == shapeDims shp2
&& map patElemName pat_els1 == winp_arrs2
&& case (mlam1, mlam2) of
(Nothing, Nothing) -> True
(Just (lam1, see1), Just (lam2, see2)) ->
(see1 == see2) && equivLambda M.empty lam1 lam2
_ -> False
substInSEs :: M.Map VName VName -> [SubExp] -> [SubExp]
substInSEs vtab = map substInSE
where
substInSE (Var x)
| Just y <- M.lookup x vtab = Var y
substInSE z = z