MapNest.hs

{-# LANGUAGE TypeFamilies #-}

module Futhark.Analysis.HORep.MapNest
  ( Nesting (..),
    MapNest (..),
    depth,
    typeOf,
    params,
    inputs,
    setInputs,
    fromSOAC,
    toSOAC,
    reshape,
  )
where

import Control.Monad (replicateM)
import Data.List (find)
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.HORep.SOAC (SOAC)
import Futhark.Analysis.HORep.SOAC qualified as SOAC
import Futhark.Construct
import Futhark.IR hiding (typeOf)
import Futhark.IR.SOACS (SOACS, isMapSOAC)
import Futhark.IR.SOACS.SOAC qualified as Futhark
import Futhark.Transform.Substitute

data Nesting = Nesting
  { nestingParamNames :: [VName],
    nestingResult :: [VName],
    nestingReturnType :: [Type],
    nestingWidth :: SubExp
  }
  deriving (Eq, Ord, Show)

data MapNest = MapNest
  { mapNestWidth :: SubExp,
    mapNestLambda :: Lambda SOACS,
    mapNestNestings :: [Nesting],
    mapNestInput :: [SOAC.Input]
  }
  deriving (Show)

depth :: MapNest -> Int
depth (MapNest _ _ nests _) = 1 + length nests

typeOf :: MapNest -> [Type]
typeOf (MapNest w lam [] _) =
  map (`arrayOfRow` w) $ lambdaReturnType lam
typeOf (MapNest w _ (nest : _) _) =
  map (`arrayOfRow` w) $ nestingReturnType nest

params :: MapNest -> [VName]
params (MapNest _ lam [] _) =
  map paramName $ lambdaParams lam
params (MapNest _ _ (nest : _) _) =
  nestingParamNames nest

inputs :: MapNest -> [SOAC.Input]
inputs (MapNest _ _ _ inps) = inps

setInputs :: [SOAC.Input] -> MapNest -> MapNest
setInputs [] (MapNest w body ns _) = MapNest w body ns []
setInputs (inp : inps) (MapNest _ body ns _) = MapNest w body ns' (inp : inps)
  where
    w = arraySize 0 $ SOAC.inputType inp
    ws = drop 1 $ arrayDims $ SOAC.inputType inp
    ns' = zipWith setDepth ns ws
    setDepth n nw = n {nestingWidth = nw}

pushIntoMapLambda ::
  (MonadFreshNames m) =>
  Stms SOACS ->
  Stm SOACS ->
  m (Maybe (Stm SOACS))
pushIntoMapLambda stms (Let pat aux (Op (Futhark.Screma w inps form)))
  | Just map_lam <- Futhark.isMapSOAC form,
    not $ any ((`namesIntersect` bound_by_stms) . freeIn) inps = do
      let lam_body = lambdaBody map_lam
          map_lam' =
            map_lam {lambdaBody = lam_body {bodyStms = stms <> bodyStms lam_body}}
      form' <- Futhark.mapSOAC map_lam'
      pure $ Just $ Let pat aux (Op (Futhark.Screma w inps form'))
  where
    bound_by_stms = namesFromList $ foldMap (patNames . stmPat) stms
pushIntoMapLambda _ _ = pure Nothing

massage :: (MonadFreshNames m) => SOAC SOACS -> m (SOAC SOACS)
massage soac@(SOAC.Screma w inps form)
  | Just lam <- Futhark.isMapSOAC form,
    Just (init_stms, last_stm) <- stmsLast $ bodyStms $ lambdaBody lam,
    all (cheap . stmExp) init_stms,
    all (`notNameIn` freeIn (bodyResult (lambdaBody lam))) $
      foldMap (patNames . stmPat) init_stms = do
      maybe_last_stm <- pushIntoMapLambda init_stms last_stm
      case maybe_last_stm of
        Just last_stm' -> do
          let lam' =
                lam {lambdaBody = (lambdaBody lam) {bodyStms = oneStm last_stm'}}
          SOAC.Screma w inps <$> Futhark.mapSOAC lam'
        Nothing -> pure soac
  where
    cheap (BasicOp BinOp {}) = True
    cheap (BasicOp SubExp {}) = True
    cheap (BasicOp CmpOp {}) = True
    cheap (BasicOp ConvOp {}) = True
    cheap (BasicOp UnOp {}) = True
    cheap _ = False
massage soac = pure soac

fromSOAC' ::
  (MonadFreshNames m, LocalScope SOACS m) =>
  [Ident] ->
  SOAC SOACS ->
  m (Maybe MapNest)
fromSOAC' bound soac = do
  screma <- massage soac

  case screma of
    SOAC.Screma w inps form | Just lam <- isMapSOAC form -> do
      let bound' = bound <> map paramIdent (lambdaParams lam)

      maybenest <- case ( stmsToList $ bodyStms $ lambdaBody lam,
                          bodyResult $ lambdaBody lam
                        ) of
        ([Let pat _ e], res)
          | map resSubExp res == map Var (patNames pat) ->
              localScope (scopeOfLParams $ lambdaParams lam) $
                SOAC.fromExp e
                  >>= either
                    (pure . Left)
                    ( fmap (Right . fmap (pat,))
                        . fromSOAC' bound'
                    )
        _ ->
          pure $ Right Nothing

      case maybenest of
        -- Do we have a nested MapNest?
        Right (Just (pat, mn@(MapNest inner_w body' ns' inps'))) -> do
          (ps, inps'') <-
            unzip
              <$> fixInputs
                w
                (zip (map paramName $ lambdaParams lam) inps)
                (zip (params mn) inps')
          let n' =
                Nesting
                  { nestingParamNames = ps,
                    nestingResult = patNames pat,
                    nestingReturnType = typeOf mn,
                    nestingWidth = inner_w
                  }
          pure $ Just $ MapNest w body' (n' : ns') inps''
        -- No nested MapNest it seems.
        _ -> do
          let isBound name
                | Just param <- find ((name ==) . identName) bound =
                    Just param
                | otherwise =
                    Nothing
              boundUsedInBody =
                mapMaybe isBound $ namesToList $ freeIn lam
          newParams <- mapM (newIdent' (<> "_wasfree")) boundUsedInBody
          let subst =
                M.fromList $
                  zip (map identName boundUsedInBody) (map identName newParams)
              inps' =
                inps
                  ++ map
                    (SOAC.addTransform (SOAC.Replicate mempty $ Shape [w]) . SOAC.identInput)
                    boundUsedInBody
              lam' =
                lam
                  { lambdaBody =
                      substituteNames subst $ lambdaBody lam,
                    lambdaParams =
                      lambdaParams lam
                        ++ [Param mempty name t | Ident name t <- newParams]
                  }
          pure $ Just $ MapNest w lam' [] inps'
    _any -> pure Nothing

fromSOAC :: (MonadFreshNames m, LocalScope SOACS m) => SOAC SOACS -> m (Maybe MapNest)
fromSOAC = fromSOAC' mempty

toSOAC :: (MonadFreshNames m, HasScope SOACS m) => MapNest -> m (SOAC SOACS)
toSOAC (MapNest w lam [] inps) =
  SOAC.Screma w inps <$> Futhark.mapSOAC lam
toSOAC (MapNest w lam (Nesting npnames nres nrettype nw : ns) inps) = do
  let nparams = zipWith (Param mempty) npnames $ map SOAC.inputRowType inps
  body <- runBodyBuilder $
    localScope (scopeOfLParams nparams) $ do
      letBindNames nres
        =<< SOAC.toExp
        =<< toSOAC (MapNest nw lam ns $ map (SOAC.identInput . paramIdent) nparams)
      pure $ varsRes nres
  let outerlam =
        Lambda
          { lambdaParams = nparams,
            lambdaBody = body,
            lambdaReturnType = nrettype
          }
  SOAC.Screma w inps <$> Futhark.mapSOAC outerlam

fixInputs ::
  (MonadFreshNames m) =>
  SubExp ->
  [(VName, SOAC.Input)] ->
  [(VName, SOAC.Input)] ->
  m [(VName, SOAC.Input)]
fixInputs w ourInps = mapM inspect
  where
    isParam x (y, _) = x == y

    inspect (_, SOAC.Input ts v _)
      | Just (p, pInp) <- find (isParam v) ourInps = do
          let pInp' = SOAC.transformRows ts pInp
          p' <- newName p
          pure (p', pInp')
    inspect (param, SOAC.Input ts a t) = do
      param' <- newVName (baseName param <> "_rep")
      pure (param', SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t)

-- | Reshape a map nest. It is assumed that any validity tests have
-- already been done. Will automatically reshape the inputs
-- appropriately.
reshape :: (MonadFreshNames m) => StmAux () -> Shape -> MapNest -> m MapNest
reshape aux shape (MapNest _ map_lam _ inps) =
  descend [] $ stripDims 1 shape
  where
    w = shapeSize 0 shape
    transform p inp =
      let shape' = shape <> arrayShape p
          inp_shape = arrayShape (SOAC.inputType inp)
          tr = SOAC.Reshape aux $ reshapeAll inp_shape shape'
       in SOAC.addTransform tr inp
    inps' = zipWith transform (map paramType $ lambdaParams map_lam) inps

    descend nests nest_shape
      | shapeRank nest_shape == 0 =
          pure $ MapNest w map_lam nests inps'
      | otherwise = do
          nest_params <- mapM (newName . paramName) $ lambdaParams map_lam
          res <-
            replicateM
              (length $ lambdaReturnType map_lam)
              (newVName "mapnest_res")
          let types =
                map (`arrayOfShape` nest_shape) $ lambdaReturnType map_lam
              nest = Nesting nest_params res types (shapeSize 0 nest_shape)
          descend (nests ++ [nest]) $ stripDims 1 nest_shape