RulesWithAccs.hs

{-# 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
  { accTup1 :: [PatElem (LetDec SOACS)],
    accTup2 :: WithAccInput SOACS,
    accTup3 :: LParam SOACS,
    accTup4 :: LParam SOACS,
    accTup5 :: (VName, Certs)
  }
  deriving (Show)

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,
      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 (AccTup pat_els_cur winp par_crt par_acc (res_nm, resCerts res_se) : rec1, rec2)
groupAccsHlp _ _ _ _ _ =
  error "Unreachable case reached in groupAccsHlp!"

matchingAccTup :: AccTup -> AccTup -> Bool
matchingAccTup
  (AccTup pat_els1 (shp1, _winp_arrs1, mlam1) _ _ _)
  (AccTup _ (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

groupCommonAccs :: [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
groupCommonAccs [] tup_accs2 =
  ([], [], tup_accs2)
groupCommonAccs (tup_acc1 : tup_accs1) tup_accs2
  | (commons2, uncommons2) <- L.partition (matchingAccTup tup_acc1) tup_accs2,
    length commons2 <= 1 =
      let (rec1, rec2, rec3) = groupCommonAccs tup_accs1 uncommons2
       in if null commons2
            then (rec1, tup_acc1 : rec2, rec3)
            else ((tup_acc1, head commons2) : rec1, rec2, rec3)
groupCommonAccs _ _ =
  error "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
  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 accTup1 acc_tup1',
      winp_2' <- concatMap ((\(_, xs, _) -> xs) . accTup2) 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 . 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
        res_lam <-
          runLambdaBuilder (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
        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
      renameLamInWAccInp (shp, inps, Just (lam, se)) = do
        lam' <- renameLambda lam
        pure (shp, inps, Just (lam', se))
      renameLamInWAccInp winp = pure winp
--
tryFuseWithAccs _ _ _ =
  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 doFusionInLambda fusedSomething wacc_id dg@DepGraph {dgGraph = g}
  | not (G.gelem wacc_id g) = pure dg
  | Just (StmNode (Let pat2 aux2 (WithAcc w_inps lam0))) <- G.lab g 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 = namesFromList $ concatMap (\(_, nms, _) -> nms) w_inps
          realConsumers n =
            nubOrd $
              map fst $
                filter (\(m, e) -> m /= n && case e of Alias {} -> False; _ -> True) $
                  G.lpre g n
          -- TransNodes that directly feed wacc and are exclusively consumed by wacc.
          trans_preds = do
            (tn_id, _) <- G.lsuc g 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 = map (\(a, _, _, _) -> a) trans_preds
          trans_out_nms = namesFromList $ map (\(_, out, _, _) -> out) trans_preds
          -- The unique Screma SoacNode that feeds all TransNodes and has no
          -- other consumers besides those TransNodes.
          soac_preds = do
            (tn_id, _, _, _) <- 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 = nubOrd soac_preds
      case (trans_preds, prod_ids) of
        (_ : _, [prod_id])
          | Just (SoacNode ots1 pat1 soac@(H.Screma {}) aux1) <- G.lab g prod_id,
            ots1 == mempty,
            all ((`notNameIn` wacc_cons_nms) . H.inputArray) (H.inputs soac),
            not $ namesIntersect trans_out_nms wacc_cons_nms ->
              attempt trans_ids trans_preds prod_id pat1 aux1 soac pat2 aux2 w_inps lam0
        _ -> pure dg
  | otherwise = pure dg
  where
    attempt trans_ids trans_preds prod_id pat1 aux1 soac pat2 aux2 w_inps lam0 = do
      let trans_info = map (\(_, out, tr, inp) -> (out, tr, inp)) trans_preds
      lam' <- renameLambda <=< runLambdaBuilder (lambdaParams lam0) $ do
        soac' <- H.toExp soac
        addStm $ Let pat1 aux1 soac'
        forM_ trans_info $ \(out, tr, inp) -> do
          (tr_aux, tr_exp) <- H.transformToExp tr 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 trans_ids prod_id pat2 aux2 w_inps lam'' = do
      void $ fusedSomething (StmNode $ Let pat2 aux2 $ WithAcc w_inps 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 = prod_id : trans_ids
          new_wacc = StmNode $ Let pat2 aux2 $ WithAcc w_inps lam''
          g' = foldr G.delNode g to_remove
          (wacc_preds, _, _, wacc_succs) = G.context g wacc_id
          -- Inherit producers of the absorbed nodes that are still in g'.
          removed_succs =
            nubOrd $
              concatMap (filter ((`G.gelem` g') . fst) . G.lsuc g) to_remove
          new_succs =
            nubOrd $
              removed_succs
                ++ map (\(e, n) -> (n, e)) (filter ((`G.gelem` g') . snd) wacc_succs)
          new_preds = filter ((`G.gelem` g') . snd) wacc_preds
          g'' = G.insNode (wacc_id, new_wacc) $ G.delNode wacc_id g'
          g''' = foldr (\(e, n) gr -> G.insEdge (n, wacc_id, e) gr) g'' new_preds
          g'''' = foldr (\(n, e) gr -> G.insEdge (wacc_id, n, e) gr) g''' new_succs
      pure dg {dgGraph = g''''}

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

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

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)