content: -- | This module implements facilities for determining whether a
  -- reduction or fold can be expressed in a closed form (i.e. not as a
  -- SOAC).
  --
  -- Right now, the module can detect only trivial cases.  In the
  -- future, we would like to make it more powerful, as well as possibly
  -- also being able to analyse sequential loops.
  module Futhark.Optimise.Simplify.Rules.ClosedForm
    ( foldClosedForm,
      loopClosedForm,
    )
  where
  
  import Control.Monad
  import Data.Map.Strict qualified as M
  import Data.Maybe
  import Futhark.Construct
  import Futhark.IR
  import Futhark.Optimise.Simplify.Rule
  import Futhark.Optimise.Simplify.Rules.Simple (VarLookup)
  import Futhark.Transform.Rename
  
  {-
  Motivation:
  
    let {*[int,x_size_27] map_computed_shape_1286} = replicate(x_size_27,
                                                               all_equal_shape_1044) in
    let {*[bool,x_size_27] map_size_checks_1292} = replicate(x_size_27, x_1291) in
    let {bool all_equal_checked_1298, int all_equal_shape_1299} =
      reduceT(fn {bool, int} (bool bacc_1293, int nacc_1294, bool belm_1295,
                              int nelm_1296) =>
                let {bool tuplit_elems_1297} = bacc_1293 && belm_1295 in
                {tuplit_elems_1297, nelm_1296},
              {True, 0}, map_size_checks_1292, map_computed_shape_1286)
  -}
  
  -- | @foldClosedForm look foldfun accargs arrargs@ determines whether
  -- each of the results of @foldfun@ can be expressed in a closed form.
  foldClosedForm ::
    (BuilderOps rep) =>
    VarLookup rep ->
    Pat (LetDec rep) ->
    Lambda rep ->
    [SubExp] ->
    [VName] ->
    RuleM rep ()
  foldClosedForm look pat lam accs arrs = do
    inputsize <- arraysSize 0 <$> mapM lookupType arrs
  
    t <- case patTypes pat of
      [Prim FloatType {}] -> cannotSimplify
      [Prim t] -> pure t
      _ -> cannotSimplify
  
    closedBody <-
      checkResults
        (patNames pat)
        inputsize
        mempty
        Int64
        knownBnds
        (map paramName (lambdaParams lam))
        (lambdaBody lam)
        accs
    isEmpty <- newVName "fold_input_is_empty"
    letBindNames [isEmpty] $
      BasicOp $
        CmpOp (CmpEq int64) inputsize (intConst Int64 0)
    letBind pat
      =<< ( Match [Var isEmpty]
              <$> (pure . Case [Just $ BoolValue True] <$> resultBodyM accs)
              <*> renameBody closedBody
              <*> pure (MatchDec [primBodyType t] MatchNormal)
          )
    where
      knownBnds = determineKnownBindings look lam accs arrs
  
  -- | @loopClosedForm pat respat merge bound bodys@ determines whether
  -- the do-loop can be expressed in a closed form.
  loopClosedForm ::
    (BuilderOps rep) =>
    Pat (LetDec rep) ->
    [(FParam rep, SubExp)] ->
    Names ->
    IntType ->
    SubExp ->
    Body rep ->
    RuleM rep ()
  loopClosedForm pat merge i it bound body = do
    t <- case patTypes pat of
      [Prim FloatType {}] -> cannotSimplify
      [Prim t] -> pure t
      _ -> cannotSimplify
  
    closedBody <-
      checkResults
        mergenames
        bound
        i
        it
        knownBnds
        (map identName mergeidents)
        body
        mergeexp
    isEmpty <- newVName "bound_is_zero"
    letBindNames [isEmpty] $
      BasicOp $
        CmpOp (CmpSlt it) bound (intConst it 0)
  
    letBind pat
      =<< ( Match [Var isEmpty]
              <$> (pure . Case [Just (BoolValue True)] <$> resultBodyM mergeexp)
              <*> renameBody closedBody
              <*> pure (MatchDec [primBodyType t] MatchNormal)
          )
    where
      (mergepat, mergeexp) = unzip merge
      mergeidents = map paramIdent mergepat
      mergenames = map paramName mergepat
      knownBnds = M.fromList $ zip mergenames mergeexp
  
  checkResults ::
    (BuilderOps rep) =>
    [VName] ->
    SubExp ->
    Names ->
    IntType ->
    M.Map VName SubExp ->
    -- | Lambda-bound
    [VName] ->
    Body rep ->
    [SubExp] ->
    RuleM rep (Body rep)
  checkResults pat size untouchable it knownBnds params body accs = do
    ((), stms) <-
      collectStms $
        zipWithM_ checkResult (zip pat res) (zip accparams accs)
    mkBodyM stms $ varsRes pat
    where
      stmMap = makeBindMap body
      (accparams, _) = splitAt (length accs) params
      res = bodyResult body
  
      nonFree = boundInBody body <> namesFromList params <> untouchable
  
      checkResult (p, SubExpRes _ (Var v)) (accparam, acc)
        | Just (BasicOp (BinOp bop x y)) <- M.lookup v stmMap,
          x /= y = do
            -- One of x,y must be *this* accumulator, and the other must
            -- be something that is free in the body.
            let isThisAccum = (== Var accparam)
            (this, el) <- liftMaybe $
              case ( (asFreeSubExp x, isThisAccum y),
                     (asFreeSubExp y, isThisAccum x)
                   ) of
                ((Just free, True), _) -> Just (acc, free)
                (_, (Just free, True)) -> Just (acc, free)
                _ -> Nothing
  
            case bop of
              LogAnd ->
                letBindNames [p] $ BasicOp $ BinOp LogAnd this el
              Add t w -> do
                size' <- asIntS t size
                letBindNames [p]
                  =<< eBinOp
                    (Add t w)
                    (eSubExp this)
                    (pure $ BasicOp $ BinOp (Mul t w) el size')
              FAdd t | Just properly_typed_size <- properFloatSize t -> do
                size' <- properly_typed_size
                letBindNames [p]
                  =<< eBinOp
                    (FAdd t)
                    (eSubExp this)
                    (pure $ BasicOp $ BinOp (FMul t) el size')
              _ -> cannotSimplify -- Um... sorry.
      checkResult _ _ = cannotSimplify
  
      asFreeSubExp :: SubExp -> Maybe SubExp
      asFreeSubExp (Var v)
        | v `nameIn` nonFree = M.lookup v knownBnds
      asFreeSubExp se = Just se
  
      properFloatSize t =
        Just $
          letSubExp "converted_size" $
            BasicOp $
              ConvOp (SIToFP it t) size
  
  determineKnownBindings ::
    VarLookup rep ->
    Lambda rep ->
    [SubExp] ->
    [VName] ->
    M.Map VName SubExp
  determineKnownBindings look lam accs arrs =
    accBnds <> arrBnds
    where
      (accparams, arrparams) =
        splitAt (length accs) $ lambdaParams lam
      accBnds =
        M.fromList $
          zip (map paramName accparams) accs
      arrBnds =
        M.fromList $
          mapMaybe isReplicate $
            zip (map paramName arrparams) arrs
  
      isReplicate (p, v)
        | Just (BasicOp (Replicate (Shape (_ : _)) ve), cs) <- look v,
          cs == mempty =
            Just (p, ve)
      isReplicate _ = Nothing
  
  makeBindMap :: Body rep -> M.Map VName (Exp rep)
  makeBindMap = M.fromList . mapMaybe isSingletonStm . stmsToList . bodyStms
    where
      isSingletonStm (Let pat _ e) = case patNames pat of
        [v] -> Just (v, e)
        _ -> Nothing
  