futhark-0.15.8: src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Optimise.InPlaceLowering.LowerIntoStm
( lowerUpdateKernels
, lowerUpdate
, LowerUpdate
, DesiredUpdate (..)
) where
import Control.Monad
import Control.Monad.Writer
import Data.List (find, unzip4)
import Data.Maybe (mapMaybe)
import Data.Either
import qualified Data.Map as M
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.Prop.Aliases
import Futhark.IR.Aliases
import Futhark.IR.Kernels
import Futhark.Construct
import Futhark.Optimise.InPlaceLowering.SubstituteIndices
data DesiredUpdate dec =
DesiredUpdate { updateName :: VName -- ^ Name of result.
, updateType :: dec -- ^ Type of result.
, updateCertificates :: Certificates
, updateSource :: VName
, updateIndices :: Slice SubExp
, updateValue :: VName
}
deriving (Show)
instance Functor DesiredUpdate where
f `fmap` u = u { updateType = f $ updateType u }
updateHasValue :: VName -> DesiredUpdate dec -> Bool
updateHasValue name = (name==) . updateValue
type LowerUpdate lore m = Scope (Aliases lore)
-> Stm (Aliases lore)
-> [DesiredUpdate (LetDec (Aliases lore))]
-> Maybe (m [Stm (Aliases lore)])
lowerUpdate :: (MonadFreshNames m, Bindable lore,
LetDec lore ~ Type, CanBeAliased (Op lore)) => LowerUpdate lore m
lowerUpdate scope (Let pat aux (DoLoop ctx val form body)) updates = do
canDo <- lowerUpdateIntoLoop scope updates pat ctx val form body
Just $ do
(prebnds, postbnds, ctxpat, valpat, ctx', val', body') <- canDo
return $
prebnds ++ [certify (stmAuxCerts aux) $
mkLet ctxpat valpat $ DoLoop ctx' val' form body'] ++ postbnds
lowerUpdate _
(Let pat aux (BasicOp (SubExp (Var v))))
[DesiredUpdate bindee_nm bindee_dec cs src is val]
| patternNames pat == [src] =
let is' = fullSlice (typeOf bindee_dec) is
in Just $
return [certify (stmAuxCerts aux <> cs) $
mkLet [] [Ident bindee_nm $ typeOf bindee_dec] $
BasicOp $ Update v is' $ Var val]
lowerUpdate _ _ _ =
Nothing
lowerUpdateKernels :: MonadFreshNames m => LowerUpdate Kernels m
lowerUpdateKernels scope
(Let pat aux (Op (SegOp (SegMap lvl space ts kbody)))) updates
| all ((`elem` patternNames pat) . updateValue) updates = do
mk <- lowerUpdatesIntoSegMap scope pat updates space kbody
Just $ do
(pat', kbody', poststms) <- mk
let cs = stmAuxCerts aux <> foldMap updateCertificates updates
return $
certify cs (Let pat' aux $ Op $ SegOp $ SegMap lvl space ts kbody') :
stmsToList poststms
lowerUpdateKernels scope stm updates = lowerUpdate scope stm updates
lowerUpdatesIntoSegMap :: MonadFreshNames m =>
Scope (Aliases Kernels)
-> Pattern (Aliases Kernels)
-> [DesiredUpdate (LetDec (Aliases Kernels))]
-> SegSpace
-> KernelBody (Aliases Kernels)
-> Maybe (m (Pattern (Aliases Kernels),
KernelBody (Aliases Kernels),
Stms (Aliases Kernels)))
lowerUpdatesIntoSegMap scope pat updates kspace kbody = do
-- The updates are all-or-nothing. Being more liberal would require
-- changes to the in-place-lowering pass itself.
mk <- zipWithM onRet (patternElements pat) (kernelBodyResult kbody)
return $ do
(pes, bodystms, krets, poststms) <- unzip4 <$> sequence mk
return (Pattern [] pes,
kbody { kernelBodyStms = kernelBodyStms kbody <> mconcat bodystms
, kernelBodyResult = krets
},
mconcat poststms)
where (gtids, _dims) = unzip $ unSegSpace kspace
onRet (PatElem v v_dec) ret
| Just (DesiredUpdate bindee_nm bindee_dec _cs src slice _val) <-
find ((==v) . updateValue) updates = do
Returns _ se <- Just ret
Just $ do
let pexp = primExpFromSubExp int32
(slice', bodystms) <- flip runBinderT scope $
traverse (toSubExp "index") $
fixSlice (map (fmap pexp) slice) $
map (pexp . Var) gtids
let res_dims = arrayDims $ snd bindee_dec
ret' = WriteReturns res_dims src [(map DimFix slice', se)]
return (PatElem bindee_nm bindee_dec,
bodystms,
ret',
oneStm $ mkLet [] [Ident v $ typeOf v_dec] $
BasicOp $ Index bindee_nm slice)
onRet pe ret =
Just $ return (pe, mempty, ret, mempty)
lowerUpdateIntoLoop :: (Bindable lore, BinderOps lore,
Aliased lore, LetDec lore ~ (als, Type),
MonadFreshNames m) =>
Scope lore
-> [DesiredUpdate (LetDec lore)]
-> Pattern lore
-> [(FParam lore, SubExp)]
-> [(FParam lore, SubExp)]
-> LoopForm lore
-> Body lore
-> Maybe (m ([Stm lore],
[Stm lore],
[Ident],
[Ident],
[(FParam lore, SubExp)],
[(FParam lore, SubExp)],
Body lore))
lowerUpdateIntoLoop scope updates pat ctx val form body = do
-- Algorithm:
--
-- 0) Map each result of the loop body to a corresponding in-place
-- update, if one exists.
--
-- 1) Create new merge variables corresponding to the arrays being
-- updated; extend the pattern and the @res@ list with these,
-- and remove the parts of the result list that have a
-- corresponding in-place update.
--
-- (The creation of the new merge variable identifiers is
-- actually done at the same time as step (0)).
--
-- 2) Create in-place updates at the end of the loop body.
--
-- 3) Create index expressions that read back the values written
-- in (2). If the merge parameter corresponding to this value
-- is unique, also @copy@ this value.
--
-- 4) Update the result of the loop body to properly pass the new
-- arrays and indexed elements to the next iteration of the
-- loop.
--
-- We also check that the merge parameters we work with have
-- loop-invariant shapes.
-- Safety condition (8).
forM_ (zip val $ bodyAliases body) $ \((p, _), als) ->
guard $ not $ paramName p `nameIn` als
mk_in_place_map <- summariseLoop updates usedInBody resmap val
Just $ do
in_place_map <- mk_in_place_map
(val',prebnds,postbnds) <- mkMerges in_place_map
let (ctxpat,valpat) = mkResAndPat in_place_map
idxsubsts = indexSubstitutions in_place_map
(idxsubsts', newbnds) <- substituteIndices idxsubsts $ bodyStms body
(body_res, res_bnds) <- manipulateResult in_place_map idxsubsts'
let body' = mkBody (newbnds<>res_bnds) body_res
return (prebnds, postbnds, ctxpat, valpat, ctx, val', body')
where usedInBody = mconcat $ map expandAliases $ namesToList $ freeIn body <> freeIn form
expandAliases v = case M.lookup v scope of
Just (LetName dec) -> oneName v <> aliasesOf dec
_ -> oneName v
resmap = zip (bodyResult body) $ patternValueIdents pat
mkMerges :: (MonadFreshNames m, Bindable lore) =>
[LoopResultSummary (als, Type)]
-> m ([(Param DeclType, SubExp)], [Stm lore], [Stm lore])
mkMerges summaries = do
((origmerge, extramerge), (prebnds, postbnds)) <-
runWriterT $ partitionEithers <$> mapM mkMerge summaries
return (origmerge ++ extramerge, prebnds, postbnds)
mkMerge summary
| Just (update, mergename, mergedec) <- relatedUpdate summary = do
source <- newVName "modified_source"
let source_t = snd $ updateType update
elmident = Ident
(updateValue update)
(source_t `setArrayDims` sliceDims (updateIndices update))
tell ([mkLet [] [Ident source source_t] $ BasicOp $ Update
(updateSource update)
(fullSlice source_t $ updateIndices update) $
snd $ mergeParam summary],
[mkLet [] [elmident] $ BasicOp $ Index
(updateName update)
(fullSlice source_t $ updateIndices update)])
return $ Right (Param
mergename
(toDecl (typeOf mergedec) Unique),
Var source)
| otherwise = return $ Left $ mergeParam summary
mkResAndPat summaries =
let (origpat,extrapat) = partitionEithers $ map mkResAndPat' summaries
in (patternContextIdents pat,
origpat ++ extrapat)
mkResAndPat' summary
| Just (update, _, _) <- relatedUpdate summary =
Right (Ident (updateName update) (snd $ updateType update))
| otherwise =
Left (inPatternAs summary)
summariseLoop :: MonadFreshNames m =>
[DesiredUpdate (als, Type)]
-> Names
-> [(SubExp, Ident)]
-> [(Param DeclType, SubExp)]
-> Maybe (m [LoopResultSummary (als, Type)])
summariseLoop updates usedInBody resmap merge =
sequence <$> zipWithM summariseLoopResult resmap merge
where summariseLoopResult (se, v) (fparam, mergeinit)
| Just update <- find (updateHasValue $ identName v) updates =
if updateSource update `nameIn` usedInBody
then Nothing
else if hasLoopInvariantShape fparam then Just $ do
lowered_array <- newVName "lowered_array"
return LoopResultSummary { resultSubExp = se
, inPatternAs = v
, mergeParam = (fparam, mergeinit)
, relatedUpdate = Just (update,
lowered_array,
updateType update)
}
else Nothing
summariseLoopResult _ _ =
Nothing -- XXX: conservative; but this entire pass is going away.
hasLoopInvariantShape = all loopInvariant . arrayDims . paramType
merge_param_names = map (paramName . fst) merge
loopInvariant (Var v) = v `notElem` merge_param_names
loopInvariant Constant{} = True
data LoopResultSummary dec =
LoopResultSummary { resultSubExp :: SubExp
, inPatternAs :: Ident
, mergeParam :: (Param DeclType, SubExp)
, relatedUpdate :: Maybe (DesiredUpdate dec, VName, dec)
}
deriving (Show)
indexSubstitutions :: [LoopResultSummary dec]
-> IndexSubstitutions dec
indexSubstitutions = mapMaybe getSubstitution
where getSubstitution res = do
(DesiredUpdate _ _ cs _ is _, nm, dec) <- relatedUpdate res
let name = paramName $ fst $ mergeParam res
return (name, (cs, nm, dec, is))
manipulateResult :: (Bindable lore, MonadFreshNames m) =>
[LoopResultSummary (LetDec lore)]
-> IndexSubstitutions (LetDec lore)
-> m (Result, Stms lore)
manipulateResult summaries substs = do
let (orig_ses,updated_ses) = partitionEithers $ map unchangedRes summaries
(subst_ses, res_bnds) <- runWriterT $ zipWithM substRes updated_ses substs
return (orig_ses ++ subst_ses, stmsFromList res_bnds)
where
unchangedRes summary =
case relatedUpdate summary of
Nothing -> Left $ resultSubExp summary
Just _ -> Right $ resultSubExp summary
substRes (Var res_v) (subst_v, (_, nm, _, _))
| res_v == subst_v =
return $ Var nm
substRes res_se (_, (cs, nm, dec, is)) = do
v' <- newIdent' (++"_updated") $ Ident nm $ typeOf dec
tell [certify cs $ mkLet [] [v'] $ BasicOp $
Update nm (fullSlice (typeOf dec) is) res_se]
return $ Var $ identName v'