futhark-0.25.33: 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
( tryFuseWithAccs,
)
where
import Control.Monad
import Data.Map.Strict qualified as M
import Futhark.Construct
import Futhark.IR.SOACS hiding (SOAC (..))
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
---------------------------------------------------
--- 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