{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}

module Futhark.Optimise.ArrayShortCircuiting.DataStructs
  ( Coalesced (..),
    CoalescedKind (..),
    ArrayMemBound (..),
    AllocTab,
    HasMemBlock,
    ScalarTab,
    CoalsTab,
    ScopeTab,
    CoalsEntry (..),
    FreeVarSubsts,
    LmadRef,
    MemRefs (..),
    AccessSummary (..),
    BotUpEnv (..),
    InhibitTab,
    unionCoalsEntry,
    vnameToPrimExp,
    getArrMemAssocFParam,
    getScopeMemInfo,
    createsNewArrOK,
    getArrMemAssoc,
    getUniqueMemFParam,
    markFailedCoal,
    accessSubtract,
    markSuccessCoal,
  )
where

import Control.Applicative
import Data.Functor ((<&>))
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.IR.Aliases
import Futhark.IR.GPUMem as GPU
import Futhark.IR.MCMem as MC
import Futhark.IR.Mem.LMAD qualified as LMAD
import Futhark.IR.SeqMem
import Futhark.Util.Pretty hiding (line, sep, (</>))
import Prelude

-- | maps array-variable names to various info, including
--   types, memory block and index function, etc.
type ScopeTab rep = Scope (Aliases rep)

-- | An LMAD specialized to TPrimExps (a typed primexp)
type LmadRef = LMAD.LMAD (TPrimExp Int64 VName)

-- | Summary of all memory accesses at a given point in the code
data AccessSummary
  = -- | The access summary was statically undeterminable, for instance by
    -- having multiple lmads. In this case, we should conservatively avoid all
    -- coalescing.
    Undeterminable
  | -- | A conservative estimate of the set of accesses up until this point.
    Set (S.Set LmadRef)

instance Semigroup AccessSummary where
  AccessSummary
Undeterminable <> :: AccessSummary -> AccessSummary -> AccessSummary
<> AccessSummary
_ = AccessSummary
Undeterminable
  AccessSummary
_ <> AccessSummary
Undeterminable = AccessSummary
Undeterminable
  (Set Set LMAD
a) <> (Set Set LMAD
b) =
    Set LMAD -> AccessSummary
Set (Set LMAD -> AccessSummary) -> Set LMAD -> AccessSummary
forall a b. (a -> b) -> a -> b
$ Set LMAD -> Set LMAD -> Set LMAD
forall a. Ord a => Set a -> Set a -> Set a
S.union Set LMAD
a Set LMAD
b

instance Monoid AccessSummary where
  mempty :: AccessSummary
mempty = Set LMAD -> AccessSummary
Set Set LMAD
forall a. Monoid a => a
mempty

instance FreeIn AccessSummary where
  freeIn' :: AccessSummary -> FV
freeIn' AccessSummary
Undeterminable = FV
forall a. Monoid a => a
mempty
  freeIn' (Set Set LMAD
s) = Set LMAD -> FV
forall a. FreeIn a => a -> FV
freeIn' Set LMAD
s

accessSubtract :: AccessSummary -> AccessSummary -> AccessSummary
accessSubtract :: AccessSummary -> AccessSummary -> AccessSummary
accessSubtract AccessSummary
Undeterminable AccessSummary
_ = AccessSummary
Undeterminable
accessSubtract AccessSummary
_ AccessSummary
Undeterminable = AccessSummary
Undeterminable
accessSubtract (Set Set LMAD
s1) (Set Set LMAD
s2) = Set LMAD -> AccessSummary
Set (Set LMAD -> AccessSummary) -> Set LMAD -> AccessSummary
forall a b. (a -> b) -> a -> b
$ Set LMAD
s1 Set LMAD -> Set LMAD -> Set LMAD
forall a. Ord a => Set a -> Set a -> Set a
S.\\ Set LMAD
s2

data MemRefs = MemRefs
  { -- | The access summary of all references (reads
    -- and writes) to the destination of a coalescing entry
    MemRefs -> AccessSummary
dstrefs :: AccessSummary,
    -- | The access summary of all writes to the source of a coalescing entry
    MemRefs -> AccessSummary
srcwrts :: AccessSummary
  }

instance Semigroup MemRefs where
  MemRefs
m1 <> :: MemRefs -> MemRefs -> MemRefs
<> MemRefs
m2 =
    AccessSummary -> AccessSummary -> MemRefs
MemRefs (MemRefs -> AccessSummary
dstrefs MemRefs
m1 AccessSummary -> AccessSummary -> AccessSummary
forall a. Semigroup a => a -> a -> a
<> MemRefs -> AccessSummary
dstrefs MemRefs
m2) (MemRefs -> AccessSummary
srcwrts MemRefs
m1 AccessSummary -> AccessSummary -> AccessSummary
forall a. Semigroup a => a -> a -> a
<> MemRefs -> AccessSummary
srcwrts MemRefs
m2)

instance Monoid MemRefs where
  mempty :: MemRefs
mempty = AccessSummary -> AccessSummary -> MemRefs
MemRefs AccessSummary
forall a. Monoid a => a
mempty AccessSummary
forall a. Monoid a => a
mempty

data CoalescedKind
  = -- | let x    = copy b^{lu}
    CopyCoal
  | -- | let x[i] = b^{lu}
    InPlaceCoal
  | -- | let x    = concat(a, b^{lu})
    ConcatCoal
  | -- | transitive, i.e., other variables aliased with b.
    TransitiveCoal
  | MapCoal

-- | Information about a memory block: type, shape, name and ixfun.
data ArrayMemBound = MemBlock
  { ArrayMemBound -> PrimType
primType :: PrimType,
    ArrayMemBound -> Shape
shape :: Shape,
    ArrayMemBound -> VName
memName :: VName,
    ArrayMemBound -> LMAD
ixfun :: LMAD
  }

-- | Free variable substitutions
type FreeVarSubsts = M.Map VName (TPrimExp Int64 VName)

-- | Coalesced Access Entry
data Coalesced
  = Coalesced
      -- | the kind of coalescing
      CoalescedKind
      -- | destination mem_block info @f_m_x[i]@ (must be ArrayMem)
      -- (Maybe IxFun) -- the inverse ixfun of a coalesced array, such that
      --                     --  ixfuns can be correctly constructed for aliases;
      ArrayMemBound
      -- | substitutions for free vars in index function
      FreeVarSubsts

data CoalsEntry = CoalsEntry
  { -- | destination memory block
    CoalsEntry -> VName
dstmem :: VName,
    -- | index function of the destination (used for rebasing)
    CoalsEntry -> LMAD
dstind :: LMAD,
    -- | aliased destination memory blocks can appear
    --   due to repeated (optimistic) coalescing.
    CoalsEntry -> Names
alsmem :: Names,
    -- | per variable-name coalesced entries
    CoalsEntry -> Map VName Coalesced
vartab :: M.Map VName Coalesced,
    -- | keys are variable names, values are memblock names;
    --   it records optimistically added coalesced nodes, e.g.,
    --   in the case of if-then-else expressions. For example:
    --       @x    = map f a@
    --       @.. use of y ..@
    --       @b    = map g a@
    --       @x[i] = b      @
    --       @y[k] = x      @
    --   the coalescing of @b@ in @x[i]@ succeeds, but
    --   is dependent of the success of the coalescing
    --   of @x@ in @y[k]@, which fails in this case
    --   because @y@ is used before the new array creation
    --   of @x = map f@. Hence @optdeps@ of the @m_b@ CoalsEntry
    --   records @x -> m_x@ and at the end of analysis it is removed
    --   from the successfully coalesced table if @m_x@ is
    --   unsuccessful.
    --   Storing @m_x@ would probably be sufficient if memory would
    --     not be reused--e.g., by register allocation on arrays--the
    --     @x@ discriminates between memory being reused across semantically
    --     different arrays (searched in @vartab@ field).
    CoalsEntry -> Map VName VName
optdeps :: M.Map VName VName,
    -- | Access summaries of uses and writes of destination and source
    -- respectively.
    CoalsEntry -> MemRefs
memrefs :: MemRefs,
    -- | Certificates of the destination, which must be propagated to
    -- the source. When short-circuiting reaches the array creation
    -- point, we must check whether the certs are in scope for
    -- short-circuiting to succeed.
    CoalsEntry -> Certs
certs :: Certs
  }

-- | the allocatted memory blocks
type AllocTab = M.Map VName Space

-- | maps a variable name to its PrimExp scalar expression
type ScalarTab = M.Map VName (PrimExp VName)

-- | maps a memory-block name to a 'CoalsEntry'. Among other things, it contains
--   @vartab@, a map in which each variable associated to that memory block is
--   bound to its 'Coalesced' info.
type CoalsTab = M.Map VName CoalsEntry

-- | inhibited memory-block mergings from the key (memory block)
--   to the value (set of memory blocks).
type InhibitTab = M.Map VName Names

data BotUpEnv = BotUpEnv
  { -- | maps scalar variables to theirs PrimExp expansion
    BotUpEnv -> ScalarTab
scals :: ScalarTab,
    -- | Optimistic coalescing info. We are currently trying to coalesce these
    -- memory blocks.
    BotUpEnv -> CoalsTab
activeCoals :: CoalsTab,
    -- | Committed (successfull) coalescing info. These memory blocks have been
    -- successfully coalesced.
    BotUpEnv -> CoalsTab
successCoals :: CoalsTab,
    -- | The coalescing failures from this pass. We will no longer try to merge
    -- these memory blocks.
    BotUpEnv -> InhibitTab
inhibit :: InhibitTab
  }

instance Pretty CoalsTab where
  pretty :: forall ann. CoalsTab -> Doc ann
pretty = [(VName, CoalsEntry)] -> Doc ann
forall ann. [(VName, CoalsEntry)] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ([(VName, CoalsEntry)] -> Doc ann)
-> (CoalsTab -> [(VName, CoalsEntry)]) -> CoalsTab -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoalsTab -> [(VName, CoalsEntry)]
forall k a. Map k a -> [(k, a)]
M.toList

instance Pretty AccessSummary where
  pretty :: forall ann. AccessSummary -> Doc ann
pretty AccessSummary
Undeterminable = Doc ann
"Undeterminable"
  pretty (Set Set LMAD
a) = Doc ann
"Access-Set:" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [LMAD] -> Doc ann
forall ann. [LMAD] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Set LMAD -> [LMAD]
forall a. Set a -> [a]
S.toList Set LMAD
a) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
" "

instance Pretty MemRefs where
  pretty :: forall ann. MemRefs -> Doc ann
pretty (MemRefs AccessSummary
a AccessSummary
b) = Doc ann
"( Use-Sum:" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> AccessSummary -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. AccessSummary -> Doc ann
pretty AccessSummary
a Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"Write-Sum:" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> AccessSummary -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. AccessSummary -> Doc ann
pretty AccessSummary
b Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
")"

instance Pretty CoalescedKind where
  pretty :: forall ann. CoalescedKind -> Doc ann
pretty CoalescedKind
CopyCoal = Doc ann
"Copy"
  pretty CoalescedKind
InPlaceCoal = Doc ann
"InPlace"
  pretty CoalescedKind
ConcatCoal = Doc ann
"Concat"
  pretty CoalescedKind
TransitiveCoal = Doc ann
"Transitive"
  pretty CoalescedKind
MapCoal = Doc ann
"Map"

instance Pretty ArrayMemBound where
  pretty :: forall ann. ArrayMemBound -> Doc ann
pretty (MemBlock PrimType
ptp Shape
shp VName
m_nm LMAD
ixfn) =
    Doc ann
"{" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. PrimType -> Doc ann
pretty PrimType
ptp Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"," Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Shape -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Shape -> Doc ann
pretty Shape
shp Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"," Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
m_nm Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"," Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> LMAD -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. LMAD -> Doc ann
pretty LMAD
ixfn Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"}"

instance Pretty Coalesced where
  pretty :: forall ann. Coalesced -> Doc ann
pretty (Coalesced CoalescedKind
knd ArrayMemBound
mbd FreeVarSubsts
_) =
    Doc ann
"(Kind:"
      Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> CoalescedKind -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. CoalescedKind -> Doc ann
pretty CoalescedKind
knd
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
", membds:"
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> ArrayMemBound -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. ArrayMemBound -> Doc ann
pretty ArrayMemBound
mbd -- <> ", subs:" <+> pretty subs
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
")"
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"\n"

instance Pretty CoalsEntry where
  pretty :: forall ann. CoalsEntry -> Doc ann
pretty CoalsEntry
etry =
    Doc ann
"{"
      Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"Dstmem:"
      Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty (CoalsEntry -> VName
dstmem CoalsEntry
etry)
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
", AliasMems:"
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Names -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Names -> Doc ann
pretty (CoalsEntry -> Names
alsmem CoalsEntry
etry)
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
", optdeps:"
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [(VName, VName)] -> Doc ann
forall ann. [(VName, VName)] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Map VName VName -> [(VName, VName)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName VName -> [(VName, VName)])
-> Map VName VName -> [(VName, VName)]
forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName VName
optdeps CoalsEntry
etry)
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
", memrefs:"
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> MemRefs -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. MemRefs -> Doc ann
pretty (CoalsEntry -> MemRefs
memrefs CoalsEntry
etry)
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
", vartab:"
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [(VName, Coalesced)] -> Doc ann
forall ann. [(VName, Coalesced)] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Map VName Coalesced -> [(VName, Coalesced)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName Coalesced -> [(VName, Coalesced)])
-> Map VName Coalesced -> [(VName, Coalesced)]
forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
etry)
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"}"
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"\n"

-- | Compute the union of two 'CoalsEntry'. If two 'CoalsEntry' do not refer to
-- the same destination memory and use the same index function, the first
-- 'CoalsEntry' is returned.
unionCoalsEntry :: CoalsEntry -> CoalsEntry -> CoalsEntry
unionCoalsEntry :: CoalsEntry -> CoalsEntry -> CoalsEntry
unionCoalsEntry CoalsEntry
etry1 (CoalsEntry VName
dstmem2 LMAD
dstind2 Names
alsmem2 Map VName Coalesced
vartab2 Map VName VName
optdeps2 MemRefs
memrefs2 Certs
certs2) =
  if CoalsEntry -> VName
dstmem CoalsEntry
etry1 VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
dstmem2 Bool -> Bool -> Bool
|| CoalsEntry -> LMAD
dstind CoalsEntry
etry1 LMAD -> LMAD -> Bool
forall a. Eq a => a -> a -> Bool
/= LMAD
dstind2
    then CoalsEntry
etry1
    else
      CoalsEntry
etry1
        { alsmem = alsmem etry1 <> alsmem2,
          optdeps = optdeps etry1 <> optdeps2,
          vartab = vartab etry1 <> vartab2,
          memrefs = memrefs etry1 <> memrefs2,
          certs = certs etry1 <> certs2
        }

-- | Get the names of array 'PatElem's in a 'Pat' and the corresponding
-- 'ArrayMemBound' information for each array.
getArrMemAssoc :: Pat (aliases, LetDecMem) -> [(VName, ArrayMemBound)]
getArrMemAssoc :: forall aliases.
Pat (aliases, LetDecMem) -> [(VName, ArrayMemBound)]
getArrMemAssoc Pat (aliases, LetDecMem)
pat =
  (PatElem (aliases, LetDecMem) -> Maybe (VName, ArrayMemBound))
-> [PatElem (aliases, LetDecMem)] -> [(VName, ArrayMemBound)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
    ( \PatElem (aliases, LetDecMem)
patel -> case (aliases, LetDecMem) -> LetDecMem
forall a b. (a, b) -> b
snd ((aliases, LetDecMem) -> LetDecMem)
-> (aliases, LetDecMem) -> LetDecMem
forall a b. (a -> b) -> a -> b
$ PatElem (aliases, LetDecMem) -> (aliases, LetDecMem)
forall dec. PatElem dec -> dec
patElemDec PatElem (aliases, LetDecMem)
patel of
        (MemArray PrimType
tp Shape
shp NoUniqueness
_ (ArrayIn VName
mem_nm LMAD
indfun)) ->
          (VName, ArrayMemBound) -> Maybe (VName, ArrayMemBound)
forall a. a -> Maybe a
Just (PatElem (aliases, LetDecMem) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (aliases, LetDecMem)
patel, PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
mem_nm LMAD
indfun)
        MemMem Space
_ -> Maybe (VName, ArrayMemBound)
forall a. Maybe a
Nothing
        MemPrim PrimType
_ -> Maybe (VName, ArrayMemBound)
forall a. Maybe a
Nothing
        MemAcc {} -> Maybe (VName, ArrayMemBound)
forall a. Maybe a
Nothing
    )
    ([PatElem (aliases, LetDecMem)] -> [(VName, ArrayMemBound)])
-> [PatElem (aliases, LetDecMem)] -> [(VName, ArrayMemBound)]
forall a b. (a -> b) -> a -> b
$ Pat (aliases, LetDecMem) -> [PatElem (aliases, LetDecMem)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (aliases, LetDecMem)
pat

-- | Get the names of arrays in a list of 'FParam' and the corresponding
-- 'ArrayMemBound' information for each array.
getArrMemAssocFParam :: [Param FParamMem] -> [(VName, Uniqueness, ArrayMemBound)]
getArrMemAssocFParam :: [Param FParamMem] -> [(VName, Uniqueness, ArrayMemBound)]
getArrMemAssocFParam =
  (Param FParamMem -> Maybe (VName, Uniqueness, ArrayMemBound))
-> [Param FParamMem] -> [(VName, Uniqueness, ArrayMemBound)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
    ( \Param FParamMem
param -> case Param FParamMem -> FParamMem
forall dec. Param dec -> dec
paramDec Param FParamMem
param of
        (MemArray PrimType
tp Shape
shp Uniqueness
u (ArrayIn VName
mem_nm LMAD
indfun)) ->
          (VName, Uniqueness, ArrayMemBound)
-> Maybe (VName, Uniqueness, ArrayMemBound)
forall a. a -> Maybe a
Just (Param FParamMem -> VName
forall dec. Param dec -> VName
paramName Param FParamMem
param, Uniqueness
u, PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
mem_nm LMAD
indfun)
        MemMem Space
_ -> Maybe (VName, Uniqueness, ArrayMemBound)
forall a. Maybe a
Nothing
        MemPrim PrimType
_ -> Maybe (VName, Uniqueness, ArrayMemBound)
forall a. Maybe a
Nothing
        MemAcc {} -> Maybe (VName, Uniqueness, ArrayMemBound)
forall a. Maybe a
Nothing
    )

-- | Get memory blocks in a list of 'FParam' that are used for unique arrays in
-- the same list of 'FParam'.
getUniqueMemFParam :: [Param FParamMem] -> M.Map VName Space
getUniqueMemFParam :: [Param FParamMem] -> Map VName Space
getUniqueMemFParam [Param FParamMem]
params =
  let mems :: Map VName Space
mems = [(VName, Space)] -> Map VName Space
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Space)] -> Map VName Space)
-> [(VName, Space)] -> Map VName Space
forall a b. (a -> b) -> a -> b
$ (Param FParamMem -> Maybe (VName, Space))
-> [Param FParamMem] -> [(VName, Space)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Param FParamMem -> Maybe (VName, Space)
forall {d} {u} {ret}.
Param (MemInfo d u ret) -> Maybe (VName, Space)
justMem [Param FParamMem]
params
      arrayMems :: Set VName
arrayMems = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (Param FParamMem -> Maybe VName) -> [Param FParamMem] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (FParamMem -> Maybe VName
forall {d}. MemInfo d Uniqueness MemBind -> Maybe VName
justArrayMem (FParamMem -> Maybe VName)
-> (Param FParamMem -> FParamMem) -> Param FParamMem -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param FParamMem -> FParamMem
forall dec. Param dec -> dec
paramDec) [Param FParamMem]
params
   in Map VName Space
mems Map VName Space -> Set VName -> Map VName Space
forall k a. Ord k => Map k a -> Set k -> Map k a
`M.restrictKeys` Set VName
arrayMems
  where
    justMem :: Param (MemInfo d u ret) -> Maybe (VName, Space)
justMem (Param Attrs
_ VName
nm (MemMem Space
sp)) = (VName, Space) -> Maybe (VName, Space)
forall a. a -> Maybe a
Just (VName
nm, Space
sp)
    justMem Param (MemInfo d u ret)
_ = Maybe (VName, Space)
forall a. Maybe a
Nothing
    justArrayMem :: MemInfo d Uniqueness MemBind -> Maybe VName
justArrayMem (MemArray PrimType
_ ShapeBase d
_ Uniqueness
Unique (ArrayIn VName
mem_nm LMAD
_)) = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
mem_nm
    justArrayMem MemInfo d Uniqueness MemBind
_ = Maybe VName
forall a. Maybe a
Nothing

class HasMemBlock rep where
  -- | Looks up 'VName' in the given scope. If it is a 'MemArray', return the
  -- 'ArrayMemBound' information for the array.
  getScopeMemInfo :: VName -> Scope rep -> Maybe ArrayMemBound

instance HasMemBlock (Aliases SeqMem) where
  getScopeMemInfo :: VName -> Scope (Aliases SeqMem) -> Maybe ArrayMemBound
getScopeMemInfo VName
r Scope (Aliases SeqMem)
scope_env0 =
    case VName
-> Scope (Aliases SeqMem) -> Maybe (NameInfo (Aliases SeqMem))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
r Scope (Aliases SeqMem)
scope_env0 of
      Just (LetName (VarAliases
_, MemArray PrimType
tp Shape
shp NoUniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Just (FParamName (MemArray PrimType
tp Shape
shp Uniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Just (LParamName (MemArray PrimType
tp Shape
shp NoUniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Maybe (NameInfo (Aliases SeqMem))
_ -> Maybe ArrayMemBound
forall a. Maybe a
Nothing

instance HasMemBlock (Aliases GPUMem) where
  getScopeMemInfo :: VName -> Scope (Aliases GPUMem) -> Maybe ArrayMemBound
getScopeMemInfo VName
r Scope (Aliases GPUMem)
scope_env0 =
    case VName
-> Scope (Aliases GPUMem) -> Maybe (NameInfo (Aliases GPUMem))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
r Scope (Aliases GPUMem)
scope_env0 of
      Just (LetName (VarAliases
_, MemArray PrimType
tp Shape
shp NoUniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Just (FParamName (MemArray PrimType
tp Shape
shp Uniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Just (LParamName (MemArray PrimType
tp Shape
shp NoUniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Maybe (NameInfo (Aliases GPUMem))
_ -> Maybe ArrayMemBound
forall a. Maybe a
Nothing

instance HasMemBlock (Aliases MCMem) where
  getScopeMemInfo :: VName -> Scope (Aliases MCMem) -> Maybe ArrayMemBound
getScopeMemInfo VName
r Scope (Aliases MCMem)
scope_env0 =
    case VName -> Scope (Aliases MCMem) -> Maybe (NameInfo (Aliases MCMem))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
r Scope (Aliases MCMem)
scope_env0 of
      Just (LetName (VarAliases
_, MemArray PrimType
tp Shape
shp NoUniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Just (FParamName (MemArray PrimType
tp Shape
shp Uniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Just (LParamName (MemArray PrimType
tp Shape
shp NoUniqueness
_ (ArrayIn VName
m LMAD
idx))) -> ArrayMemBound -> Maybe ArrayMemBound
forall a. a -> Maybe a
Just (PrimType -> Shape -> VName -> LMAD -> ArrayMemBound
MemBlock PrimType
tp Shape
shp VName
m LMAD
idx)
      Maybe (NameInfo (Aliases MCMem))
_ -> Maybe ArrayMemBound
forall a. Maybe a
Nothing

-- | @True@ if the expression returns a "fresh" array.
createsNewArrOK :: Exp rep -> Bool
createsNewArrOK :: forall rep. Exp rep -> Bool
createsNewArrOK (BasicOp Replicate {}) = Bool
True
createsNewArrOK (BasicOp Iota {}) = Bool
True
createsNewArrOK (BasicOp Manifest {}) = Bool
True
createsNewArrOK (BasicOp Concat {}) = Bool
True
createsNewArrOK (BasicOp ArrayLit {}) = Bool
True
createsNewArrOK (BasicOp ArrayVal {}) = Bool
True
createsNewArrOK (BasicOp Scratch {}) = Bool
True
createsNewArrOK Exp rep
_ = Bool
False

-- | Memory-block removal from active-coalescing table
--   should only be handled via this function, it is easy
--   to run into infinite execution problem; i.e., the
--   fix-pointed iteration of coalescing transformation
--   assumes that whenever a coalescing fails it is
--   recorded in the @inhibit@ table.
markFailedCoal ::
  (CoalsTab, InhibitTab) ->
  VName ->
  (CoalsTab, InhibitTab)
markFailedCoal :: (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
coal_tab, InhibitTab
inhb_tab) VName
src_mem =
  case VName -> CoalsTab -> Maybe CoalsEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
src_mem CoalsTab
coal_tab of
    Maybe CoalsEntry
Nothing -> (CoalsTab
coal_tab, InhibitTab
inhb_tab)
    Just CoalsEntry
coale ->
      let failed_set :: Names
failed_set = VName -> Names
oneName (VName -> Names) -> VName -> Names
forall a b. (a -> b) -> a -> b
$ CoalsEntry -> VName
dstmem CoalsEntry
coale
          failed_set' :: Names
failed_set' = Names
failed_set Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names -> Maybe Names -> Names
forall a. a -> Maybe a -> a
fromMaybe Names
forall a. Monoid a => a
mempty (VName -> InhibitTab -> Maybe Names
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
src_mem InhibitTab
inhb_tab)
       in ( VName -> CoalsTab -> CoalsTab
forall k a. Ord k => k -> Map k a -> Map k a
M.delete VName
src_mem CoalsTab
coal_tab,
            VName -> Names -> InhibitTab -> InhibitTab
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
src_mem Names
failed_set' InhibitTab
inhb_tab
          )

-- | promotion from active-to-successful coalescing tables
--   should be handled with this function (for clarity).
markSuccessCoal ::
  (CoalsTab, CoalsTab) ->
  VName ->
  CoalsEntry ->
  (CoalsTab, CoalsTab)
markSuccessCoal :: (CoalsTab, CoalsTab) -> VName -> CoalsEntry -> (CoalsTab, CoalsTab)
markSuccessCoal (CoalsTab
actv, CoalsTab
succc) VName
m_b CoalsEntry
info_b =
  ( VName -> CoalsTab -> CoalsTab
forall k a. Ord k => k -> Map k a -> Map k a
M.delete VName
m_b CoalsTab
actv,
    VName -> CoalsEntry -> CoalsTab -> CoalsTab
appendCoalsInfo VName
m_b CoalsEntry
info_b CoalsTab
succc
  )

-- | merges entries in the coalesced table.
appendCoalsInfo :: VName -> CoalsEntry -> CoalsTab -> CoalsTab
appendCoalsInfo :: VName -> CoalsEntry -> CoalsTab -> CoalsTab
appendCoalsInfo VName
mb CoalsEntry
info_new CoalsTab
coalstab =
  case VName -> CoalsTab -> Maybe CoalsEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
mb CoalsTab
coalstab of
    Maybe CoalsEntry
Nothing -> VName -> CoalsEntry -> CoalsTab -> CoalsTab
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
mb CoalsEntry
info_new CoalsTab
coalstab
    Just CoalsEntry
info_old -> VName -> CoalsEntry -> CoalsTab -> CoalsTab
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
mb (CoalsEntry -> CoalsEntry -> CoalsEntry
unionCoalsEntry CoalsEntry
info_old CoalsEntry
info_new) CoalsTab
coalstab

-- | Attempt to convert a 'VName' to a PrimExp.
--
-- First look in 'ScalarTab' to see if we have recorded the scalar value of the
-- argument. Otherwise look up the type of the argument and return a 'LeafExp'
-- if it is a 'PrimType'.
vnameToPrimExp ::
  (AliasableRep rep) =>
  ScopeTab rep ->
  ScalarTab ->
  VName ->
  Maybe (PrimExp VName)
vnameToPrimExp :: forall rep.
AliasableRep rep =>
ScopeTab rep -> ScalarTab -> VName -> Maybe (PrimExp VName)
vnameToPrimExp ScopeTab rep
scopetab ScalarTab
scaltab VName
v =
  VName -> ScalarTab -> Maybe (PrimExp VName)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v ScalarTab
scaltab
    Maybe (PrimExp VName)
-> Maybe (PrimExp VName) -> Maybe (PrimExp VName)
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ( VName -> ScopeTab rep -> Maybe (NameInfo (Aliases rep))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v ScopeTab rep
scopetab
            Maybe (NameInfo (Aliases rep))
-> (NameInfo (Aliases rep) -> Maybe PrimType) -> Maybe PrimType
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= TypeBase Shape NoUniqueness -> Maybe PrimType
forall shp u. TypeBase shp u -> Maybe PrimType
toPrimType (TypeBase Shape NoUniqueness -> Maybe PrimType)
-> (NameInfo (Aliases rep) -> TypeBase Shape NoUniqueness)
-> NameInfo (Aliases rep)
-> Maybe PrimType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NameInfo (Aliases rep) -> TypeBase Shape NoUniqueness
forall t. Typed t => t -> TypeBase Shape NoUniqueness
typeOf
            Maybe PrimType
-> (PrimType -> PrimExp VName) -> Maybe (PrimExp VName)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
v
        )

-- | Attempt to extract the 'PrimType' from a 'TypeBase'.
toPrimType :: TypeBase shp u -> Maybe PrimType
toPrimType :: forall shp u. TypeBase shp u -> Maybe PrimType
toPrimType (Prim PrimType
pt) = PrimType -> Maybe PrimType
forall a. a -> Maybe a
Just PrimType
pt
toPrimType TypeBase shp u
_ = Maybe PrimType
forall a. Maybe a
Nothing