futhark-0.22.4: src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Optimise.ArrayShortCircuiting.MemRefAggreg
( recordMemRefUses,
freeVarSubstitutions,
translateAccessSummary,
aggSummaryLoopTotal,
aggSummaryLoopPartial,
aggSummaryMapPartial,
aggSummaryMapTotal,
noMemOverlap,
)
where
import Control.Monad
import Data.Function ((&))
import Data.List (intersect, partition, uncons)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.Analysis.AlgSimplify
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.Aliases
import Futhark.IR.Mem
import Futhark.IR.Mem.IxFun qualified as IxFun
import Futhark.MonadFreshNames
import Futhark.Optimise.ArrayShortCircuiting.DataStructs
import Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis
import Futhark.Util
-----------------------------------------------------
-- Some translations of Accesses and Ixfuns --
-----------------------------------------------------
-- | Checks whether the index function can be translated at the current program
-- point and also returns the substitutions. It comes down to answering the
-- question: "can one perform enough substitutions (from the bottom-up scalar
-- table) until all vars appearing in the index function are defined in the
-- current scope?"
freeVarSubstitutions ::
FreeIn a =>
ScopeTab rep ->
ScalarTab ->
a ->
Maybe FreeVarSubsts
freeVarSubstitutions scope0 scals0 indfun =
freeVarSubstitutions' mempty $ namesToList $ freeIn indfun
where
freeVarSubstitutions' :: FreeVarSubsts -> [VName] -> Maybe FreeVarSubsts
freeVarSubstitutions' subs [] = Just subs
freeVarSubstitutions' subs0 fvs =
let fvs_not_in_scope = filter (`M.notMember` scope0) fvs
in case unzip <$> mapM getSubstitution fvs_not_in_scope of
-- We require that all free variables can be substituted
Just (subs, new_fvs) ->
freeVarSubstitutions' (subs0 <> mconcat subs) $ concat new_fvs
Nothing -> Nothing
getSubstitution v
| Just pe <- M.lookup v scals0,
IntType _ <- primExpType pe =
Just (M.singleton v $ TPrimExp pe, namesToList $ freeIn pe)
getSubstitution _v = Nothing
-- | Translates free variables in an access summary
translateAccessSummary :: ScopeTab rep -> ScalarTab -> AccessSummary -> AccessSummary
translateAccessSummary _ _ Undeterminable = Undeterminable
translateAccessSummary scope0 scals0 (Set slmads)
| Just subs <- freeVarSubstitutions scope0 scals0 slmads =
slmads
& S.map (IxFun.substituteInLMAD subs)
& Set
translateAccessSummary _ _ _ = Undeterminable
-- | This function computes the written and read memory references for the current statement
getUseSumFromStm ::
(Op rep ~ MemOp inner, HasMemBlock (Aliases rep)) =>
TopdownEnv rep ->
CoalsTab ->
Stm (Aliases rep) ->
-- | A pair of written and written+read memory locations, along with their
-- associated array and the index function used
Maybe ([(VName, VName, IxFun)], [(VName, VName, IxFun)])
getUseSumFromStm td_env coal_tab (Let _ _ (BasicOp (Index arr (Slice slc))))
| Just (MemBlock _ shp _ _) <- getScopeMemInfo arr (scope td_env),
length slc == length (shapeDims shp) && all isFix slc = do
(mem_b, mem_arr, ixfn_arr) <- getDirAliasedIxfn td_env coal_tab arr
let new_ixfn = IxFun.slice ixfn_arr $ Slice $ map (fmap pe64) slc
pure ([], [(mem_b, mem_arr, new_ixfn)])
where
isFix DimFix {} = True
isFix _ = False
getUseSumFromStm _ _ (Let Pat {} _ (BasicOp Index {})) = Just ([], []) -- incomplete slices
getUseSumFromStm _ _ (Let Pat {} _ (BasicOp FlatIndex {})) = Just ([], []) -- incomplete slices
getUseSumFromStm td_env coal_tab (Let (Pat pes) _ (BasicOp (ArrayLit ses _))) =
let rds = mapMaybe (getDirAliasedIxfn td_env coal_tab) $ mapMaybe seName ses
wrts = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) pes
in Just (wrts, wrts ++ rds)
where
seName (Var a) = Just a
seName (Constant _) = Nothing
-- In place update @x[slc] <- a@. In the "in-place update" case,
-- summaries should be added after the old variable @x@ has
-- been added in the active coalesced table.
getUseSumFromStm td_env coal_tab (Let (Pat [x']) _ (BasicOp (Update _ _x (Slice slc) a_se))) = do
(m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x')
let x_ixfn_slc = IxFun.slice x_ixfn $ Slice $ map (fmap pe64) slc
r1 = (m_b, m_x, x_ixfn_slc)
case a_se of
Constant _ -> Just ([r1], [r1])
Var a -> case getDirAliasedIxfn td_env coal_tab a of
Nothing -> Just ([r1], [r1])
Just r2 -> Just ([r1], [r1, r2])
getUseSumFromStm td_env coal_tab (Let (Pat [y]) _ (BasicOp (Copy x))) = do
-- y = copy x
wrt <- getDirAliasedIxfn td_env coal_tab $ patElemName y
rd <- getDirAliasedIxfn td_env coal_tab x
pure ([wrt], [wrt, rd])
getUseSumFromStm _ _ (Let Pat {} _ (BasicOp Copy {})) = error "Impossible"
getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Concat _i (a :| bs) _ses))) =
-- concat
let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys
rs = mapMaybe (getDirAliasedIxfn td_env coal_tab) (a : bs)
in Just (ws, ws ++ rs)
getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Manifest _perm x))) =
let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys
rs = mapMaybe (getDirAliasedIxfn td_env coal_tab) [x]
in Just (ws, ws ++ rs)
getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Replicate _shp se))) =
let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys
in case se of
Constant _ -> Just (ws, ws)
Var x -> Just (ws, ws ++ mapMaybe (getDirAliasedIxfn td_env coal_tab) [x])
getUseSumFromStm td_env coal_tab (Let (Pat [x]) _ (BasicOp (FlatUpdate _ (FlatSlice offset slc) v)))
| Just (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x) =
let x_ixfn_slc = IxFun.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc
r1 = (m_b, m_x, x_ixfn_slc)
in case getDirAliasedIxfn td_env coal_tab v of
Nothing -> Just ([r1], [r1])
Just r2 -> Just ([r1], [r1, r2])
getUseSumFromStm _ _ (Let Pat {} _ BasicOp {}) = Just ([], [])
getUseSumFromStm _ _ (Let Pat {} _ (Op (Alloc _ _))) = Just ([], [])
getUseSumFromStm _ _ _ =
-- if-then-else, loops are supposed to be treated separately,
-- calls are not supported, and Ops are not yet supported
Nothing
-- | This function:
-- 1. computes the written and read memory references for the current statement
-- (by calling @getUseSumFromStm@)
-- 2. fails the entries in active coalesced table for which the write set
-- overlaps the uses of the destination (to that point)
recordMemRefUses ::
(CanBeAliased (Op rep), RepTypes rep, Op rep ~ MemOp inner, HasMemBlock (Aliases rep)) =>
TopdownEnv rep ->
BotUpEnv ->
Stm (Aliases rep) ->
(CoalsTab, InhibitTab)
recordMemRefUses td_env bu_env stm =
let active_tab = activeCoals bu_env
inhibit_tab = inhibit bu_env
active_etries = M.toList active_tab
in case getUseSumFromStm td_env active_tab stm of
Nothing ->
M.toList active_tab
& foldl
( \state (m_b, entry) ->
if not $ null $ patNames (stmPat stm) `intersect` M.keys (vartab entry)
then markFailedCoal state m_b
else state
)
(active_tab, inhibit_tab)
Just use_sums ->
let (mb_wrts, prev_uses, mb_lmads) =
map (checkOverlapAndExpand use_sums active_tab) active_etries
& unzip3
-- keep only the entries that do not overlap with the memory
-- blocks defined in @pat@ or @inner_free_vars@.
-- the others must be recorded in @inhibit_tab@ because
-- they violate the 3rd safety condition.
active_tab1 =
M.fromList
$ map
( \(wrts, (uses, prev_use, (k, etry))) ->
let mrefs' = (memrefs etry) {dstrefs = prev_use}
etry' = etry {memrefs = mrefs'}
in (k, addLmads wrts uses etry')
)
$ mapMaybe (\(x, y) -> (,y) <$> x) -- only keep successful coals
$ zip mb_wrts
$ zip3 mb_lmads prev_uses active_etries
failed_tab =
M.fromList $
map snd $
filter (isNothing . fst) $
zip mb_wrts active_etries
(_, inhibit_tab1) = foldl markFailedCoal (failed_tab, inhibit_tab) $ M.keys failed_tab
in (active_tab1, inhibit_tab1)
where
checkOverlapAndExpand (stm_wrts, stm_uses) active_tab (m_b, etry) =
let alias_m_b = getAliases mempty m_b
stm_uses' = filter ((`notNameIn` alias_m_b) . tupFst) stm_uses
all_aliases = foldl getAliases mempty $ namesToList $ alsmem etry
ixfns = map tupThd $ filter ((`nameIn` all_aliases) . tupSnd) stm_uses'
lmads' = mapMaybe mbLmad ixfns
lmads'' =
if length lmads' == length ixfns
then Set $ S.fromList lmads'
else Undeterminable
wrt_ixfns = map tupThd $ filter ((`nameIn` alias_m_b) . tupFst) stm_wrts
wrt_tmps = mapMaybe mbLmad wrt_ixfns
prev_use =
translateAccessSummary (scope td_env) (scalarTable td_env) $
(dstrefs . memrefs) etry
wrt_lmads' =
if length wrt_tmps == length wrt_ixfns
then Set $ S.fromList wrt_tmps
else Undeterminable
original_mem_aliases =
fmap tupFst stm_uses
& uncons
& fmap fst
& (=<<) (`M.lookup` active_tab)
& maybe mempty alsmem
(wrt_lmads'', lmads) =
if m_b `nameIn` original_mem_aliases
then (wrt_lmads' <> lmads'', Set mempty)
else (wrt_lmads', lmads'')
no_overlap = noMemOverlap td_env prev_use wrt_lmads''
wrt_lmads =
if no_overlap
then Just wrt_lmads''
else Nothing
in (wrt_lmads, prev_use, lmads)
tupFst (a, _, _) = a
tupSnd (_, b, _) = b
tupThd (_, _, c) = c
getAliases acc m =
oneName m
<> acc
<> fromMaybe mempty (M.lookup m (m_alias td_env))
mbLmad indfun
| Just subs <- freeVarSubstitutions (scope td_env) (scals bu_env) indfun,
(IxFun.IxFun (lmad :| []) _ _) <- IxFun.substituteInIxFun subs indfun =
Just lmad
mbLmad _ = Nothing
addLmads wrts uses etry =
etry {memrefs = MemRefs uses wrts <> memrefs etry}
-- | Check for memory overlap of two access summaries.
--
-- This check is conservative, so unless we can guarantee that there is no
-- overlap, we return 'False'.
noMemOverlap :: (CanBeAliased (Op rep), RepTypes rep) => TopdownEnv rep -> AccessSummary -> AccessSummary -> Bool
noMemOverlap _ _ (Set mr)
| mr == mempty = True
noMemOverlap _ (Set mr) _
| mr == mempty = True
noMemOverlap td_env (Set is0) (Set js0)
| Just non_negs <- mapM (primExpFromSubExpM (vnameToPrimExp (scope td_env) (scalarTable td_env)) . Var) $ namesToList $ nonNegatives td_env =
let (_, not_disjoints) =
partition
( \i ->
all
( \j ->
IxFun.disjoint less_thans (nonNegatives td_env) i j
|| IxFun.disjoint2 () () less_thans (nonNegatives td_env) i j
|| IxFun.disjoint3 (typeOf <$> scope td_env) asserts less_thans non_negs i j
)
js
)
is
in null not_disjoints
where
less_thans = map (fmap $ fixPoint $ substituteInPrimExp $ scalarTable td_env) $ knownLessThan td_env
asserts = map (fixPoint (substituteInPrimExp $ scalarTable td_env) . primExpFromSubExp Bool) $ td_asserts td_env
is = map (fixPoint (IxFun.substituteInLMAD $ TPrimExp <$> scalarTable td_env)) $ S.toList is0
js = map (fixPoint (IxFun.substituteInLMAD $ TPrimExp <$> scalarTable td_env)) $ S.toList js0
noMemOverlap _ _ _ = False
-- | Computes the total aggregated access summary for a loop by expanding the
-- access summary given according to the iterator variable and bounds of the
-- loop.
--
-- Corresponds to:
--
-- \[
-- \bigcup_{j=0}^{j<n} Access_j
-- \]
aggSummaryLoopTotal ::
MonadFreshNames m =>
ScopeTab rep ->
ScopeTab rep ->
ScalarTab ->
Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName)) ->
AccessSummary ->
m AccessSummary
aggSummaryLoopTotal _ _ _ _ Undeterminable = pure Undeterminable
aggSummaryLoopTotal _ _ _ _ (Set l)
| l == mempty = pure $ Set mempty
aggSummaryLoopTotal scope_bef scope_loop scals_loop _ access
| Set ls <- translateAccessSummary scope_loop scals_loop access,
nms <- foldl (<>) mempty $ map freeIn $ S.toList ls,
all inBeforeScope $ namesToList nms = do
pure $ Set ls
where
inBeforeScope v =
case M.lookup v scope_bef of
Nothing -> False
Just _ -> True
aggSummaryLoopTotal _ _ scalars_loop (Just (iterator_var, (lower_bound, upper_bound))) (Set lmads) =
concatMapM
( aggSummaryOne iterator_var lower_bound upper_bound
. fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars_loop)
)
(S.toList lmads)
aggSummaryLoopTotal _ _ _ _ _ = pure Undeterminable
-- | For a given iteration of the loop $i$, computes the aggregated loop access
-- summary of all later iterations.
--
-- Corresponds to:
--
-- \[
-- \bigcup_{j=i+1}^{j<n} Access_j
-- \]
aggSummaryLoopPartial ::
MonadFreshNames m =>
ScalarTab ->
Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName)) ->
AccessSummary ->
m AccessSummary
aggSummaryLoopPartial _ _ Undeterminable = pure Undeterminable
aggSummaryLoopPartial _ Nothing _ = pure Undeterminable
aggSummaryLoopPartial scalars_loop (Just (iterator_var, (_, upper_bound))) (Set lmads) = do
-- map over each index function in the access summary
-- Substitube a fresh variable k for the loop iterator
-- if k is in stride or span of ixfun: fall back to total
-- new_stride = old_offset - old_offset (where k+1 is substituted for k)
-- new_offset = old_offset where k = lower bound of iteration
-- new_span = upper bound of iteration
concatMapM
( aggSummaryOne
iterator_var
(isInt64 (LeafExp iterator_var $ IntType Int64) + 1)
(upper_bound - typedLeafExp iterator_var - 1)
. fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars_loop)
)
(S.toList lmads)
-- | For a given map with $k$ dimensions and an index $i$ for each dimension,
-- compute the aggregated access summary of all other threads.
--
-- For the innermost dimension, this corresponds to
--
-- \[
-- \bigcup_{j=0}^{j<i} Access_j \cup \bigcup_{j=i+1}^{j<n} Access_j
-- \]
--
-- where $Access_j$ describes the point accesses in the map. As we move up in
-- dimensionality, the previous access summaries are kept, in addition to the
-- total aggregation of the inner dimensions. For outer dimensions, the equation
-- is the same, the point accesses in $Access_j$ are replaced with the total
-- aggregation of the inner dimensions.
aggSummaryMapPartial :: MonadFreshNames m => ScalarTab -> [(VName, SubExp)] -> LmadRef -> m AccessSummary
aggSummaryMapPartial _ [] = const $ pure mempty
aggSummaryMapPartial scalars dims =
helper mempty (reverse dims) . Set . S.singleton -- Reverse dims so we work from the inside out
where
helper acc [] _ = pure acc
helper Undeterminable _ _ = pure Undeterminable
helper _ _ Undeterminable = pure Undeterminable
helper (Set acc) ((gtid, size) : rest) (Set as) = do
partial_as <- aggSummaryMapPartialOne scalars (gtid, size) (Set as)
total_as <-
concatMapM
(aggSummaryOne gtid 0 (TPrimExp $ primExpFromSubExp (IntType Int64) size))
(S.toList as)
helper (Set acc <> partial_as) rest total_as
-- | Given an access summary $a$, a thread id $i$ and the size $n$ of the
-- dimension, compute the partial map summary.
--
-- Corresponds to
--
-- \[
-- \bigcup_{j=0}^{j<i} a_j \cup \bigcup_{j=i+1}^{j<n} a_j
-- \]
aggSummaryMapPartialOne :: MonadFreshNames m => ScalarTab -> (VName, SubExp) -> AccessSummary -> m AccessSummary
aggSummaryMapPartialOne _ _ Undeterminable = pure Undeterminable
aggSummaryMapPartialOne _ (_, Constant n) (Set _) | oneIsh n = pure mempty
aggSummaryMapPartialOne scalars (gtid, size) (Set lmads0) =
concatMapM
helper
[ (0, isInt64 (LeafExp gtid $ IntType Int64)),
( isInt64 (LeafExp gtid $ IntType Int64) + 1,
isInt64 (primExpFromSubExp (IntType Int64) size) - isInt64 (LeafExp gtid $ IntType Int64) - 1
)
]
where
lmads = map (fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars)) $ S.toList lmads0
helper (x, y) = concatMapM (aggSummaryOne gtid x y) lmads
-- | Computes to total access summary over a multi-dimensional map.
aggSummaryMapTotal :: MonadFreshNames m => ScalarTab -> [(VName, SubExp)] -> AccessSummary -> m AccessSummary
aggSummaryMapTotal _ [] _ = pure mempty
aggSummaryMapTotal _ _ (Set lmads)
| lmads == mempty = pure mempty
aggSummaryMapTotal _ _ Undeterminable = pure Undeterminable
aggSummaryMapTotal scalars segspace (Set lmads0) =
foldM
( \as' (gtid', size') -> case as' of
Set lmads' ->
concatMapM
( aggSummaryOne gtid' 0 $
TPrimExp $
primExpFromSubExp (IntType Int64) size'
)
(S.toList lmads')
Undeterminable -> pure Undeterminable
)
(Set lmads)
(reverse segspace)
where
lmads =
S.fromList $
map (fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars)) $
S.toList lmads0
-- | Helper function that aggregates the accesses of single LMAD according to a
-- given iterator value, a lower bound and a span.
--
-- If successful, the result is an index function with an extra outer
-- dimension. The stride of the outer dimension is computed by taking the
-- difference between two points in the index function.
--
-- The function returns 'Underterminable' if the iterator is free in the output
-- LMAD or the dimensions of the input LMAD .
aggSummaryOne :: MonadFreshNames m => VName -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> LmadRef -> m AccessSummary
aggSummaryOne iterator_var lower_bound spn lmad@(IxFun.LMAD offset0 dims0)
| iterator_var `nameIn` freeIn dims0 = pure Undeterminable
| iterator_var `notNameIn` freeIn offset0 = pure $ Set $ S.singleton lmad
| otherwise = do
new_var <- newVName "k"
let offset = replaceIteratorWith (typedLeafExp new_var) offset0
offsetp1 = replaceIteratorWith (typedLeafExp new_var + 1) offset0
new_stride = TPrimExp $ constFoldPrimExp $ simplify $ untyped $ offsetp1 - offset
new_offset = replaceIteratorWith lower_bound offset0
new_lmad =
IxFun.LMAD new_offset $
IxFun.LMADDim new_stride spn 0 IxFun.Inc : map incPerm dims0
if new_var `nameIn` freeIn new_lmad
then pure Undeterminable
else pure $ Set $ S.singleton new_lmad
where
incPerm dim = dim {IxFun.ldPerm = IxFun.ldPerm dim + 1}
replaceIteratorWith se = TPrimExp . substituteInPrimExp (M.singleton iterator_var $ untyped se) . untyped
-- | Takes a 'VName' and converts it into a 'TPrimExp' with type 'Int64'.
typedLeafExp :: VName -> TPrimExp Int64 VName
typedLeafExp vname = isInt64 $ LeafExp vname (IntType Int64)