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

module Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis
  ( TopdownEnv (..),
    ScopeTab,
    TopDownHelper,
    InhibitTab,
    updateTopdownEnv,
    updateTopdownEnvLoop,
    getDirAliasedIxfn,
    getDirAliasedIxfn',
    addInvAliassesVarTab,
    areAnyAliased,
    isInScope,
    nonNegativesInPat,
  )
where

import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.Aliases
import Futhark.IR.GPUMem as GPU
import Futhark.IR.MCMem as MC
import Futhark.IR.Mem.IxFun qualified as IxFun
import Futhark.Optimise.ArrayShortCircuiting.DataStructs

type DirAlias = IxFun -> IxFun
-- ^ A direct aliasing transformation

type InvAlias = Maybe (IxFun -> IxFun)
-- ^ An inverse aliasing transformation

type VarAliasTab = M.Map VName (VName, DirAlias, InvAlias)

type MemAliasTab = M.Map VName Names

data TopdownEnv rep = TopdownEnv
  { -- | contains the already allocated memory blocks
    forall rep. TopdownEnv rep -> AllocTab
alloc :: AllocTab,
    -- | variable info, including var-to-memblock assocs
    forall rep. TopdownEnv rep -> ScopeTab rep
scope :: ScopeTab rep,
    -- | the inherited inhibitions from the previous try
    forall rep. TopdownEnv rep -> InhibitTab
inhibited :: InhibitTab,
    -- | for statements such as transpose, reshape, index, etc., that alias
    --   an array variable: maps var-names to pair of aliased var name
    --   and index function transformation. For example, for
    --   @let b = a[slc]@ it should add the binding
    --   @ b |-> (a, `slice` slc )@
    forall rep. TopdownEnv rep -> VarAliasTab
v_alias :: VarAliasTab,
    -- | keeps track of memory block aliasing.
    --   this needs to be implemented
    forall rep. TopdownEnv rep -> InhibitTab
m_alias :: MemAliasTab,
    -- | Contains symbol information about the variables in the program. Used to
    -- determine if a variable is non-negative.
    forall rep. TopdownEnv rep -> Names
nonNegatives :: Names,
    forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
scalarTable :: M.Map VName (PrimExp VName),
    -- | A list of known relations of the form 'VName' @<@ 'SubExp', typically
    -- gotten from 'LoopForm' and 'SegSpace'.
    forall rep. TopdownEnv rep -> [(VName, PrimExp VName)]
knownLessThan :: [(VName, PrimExp VName)],
    -- | A list of the asserts encountered so far
    forall rep. TopdownEnv rep -> [SubExp]
td_asserts :: [SubExp]
  }

isInScope :: TopdownEnv rep -> VName -> Bool
isInScope :: forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env VName
m =
  VName
m forall k a. Ord k => k -> Map k a -> Bool
`M.member` forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env

-- | Get alias and (direct) index function mapping from expression
--
-- For instance, if the expression is a 'Rotate', returns the value being
-- rotated as well as a function for rotating an index function the appropriate
-- amount.
getDirAliasFromExp :: Exp (Aliases rep) -> Maybe (VName, DirAlias)
getDirAliasFromExp :: forall rep. Exp (Aliases rep) -> Maybe (VName, DirAlias)
getDirAliasFromExp (BasicOp (SubExp (Var VName
x))) = forall a. a -> Maybe a
Just (VName
x, forall a. a -> a
id)
getDirAliasFromExp (BasicOp (Opaque OpaqueOp
_ (Var VName
x))) = forall a. a -> Maybe a
Just (VName
x, forall a. a -> a
id)
getDirAliasFromExp (BasicOp (Reshape ReshapeKind
ReshapeCoerce Shape
shp VName
x)) =
  forall a. a -> Maybe a
Just (VName
x, (forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Shape num -> IxFun num
`IxFun.coerce` forall d. ShapeBase d -> [d]
shapeDims (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Shape
shp)))
getDirAliasFromExp (BasicOp (Reshape ReshapeKind
ReshapeArbitrary Shape
shp VName
x)) =
  forall a. a -> Maybe a
Just (VName
x, (forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Shape num -> IxFun num
`IxFun.reshape` forall d. ShapeBase d -> [d]
shapeDims (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Shape
shp)))
getDirAliasFromExp (BasicOp (Rearrange [Int]
_ VName
_)) =
  forall a. Maybe a
Nothing
getDirAliasFromExp (BasicOp (Rotate [SubExp]
_ VName
_)) =
  forall a. Maybe a
Nothing -- Just (x, (`IxFun.rotate` fmap pe64 rs))
getDirAliasFromExp (BasicOp (Index VName
x Slice SubExp
slc)) =
  forall a. a -> Maybe a
Just (VName
x, (forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
`IxFun.slice` (forall d. [DimIndex d] -> Slice d
Slice forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) forall a b. (a -> b) -> a -> b
$ forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slc)))
getDirAliasFromExp (BasicOp (Update Safety
_ VName
x Slice SubExp
_ SubExp
_elm)) = forall a. a -> Maybe a
Just (VName
x, forall a. a -> a
id)
getDirAliasFromExp (BasicOp (FlatIndex VName
x (FlatSlice SubExp
offset [FlatDimIndex SubExp]
idxs))) =
  forall a. a -> Maybe a
Just
    ( VName
x,
      ( forall num.
(Eq num, IntegralExp num) =>
IxFun num -> FlatSlice num -> IxFun num
`IxFun.flatSlice`
          ( forall d. d -> [FlatDimIndex d] -> FlatSlice d
FlatSlice (SubExp -> TPrimExp Int64 VName
pe64 SubExp
offset) forall a b. (a -> b) -> a -> b
$
              forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) [FlatDimIndex SubExp]
idxs
          )
      )
    )
getDirAliasFromExp (BasicOp (FlatUpdate VName
x FlatSlice SubExp
_ VName
_)) = forall a. a -> Maybe a
Just (VName
x, forall a. a -> a
id)
getDirAliasFromExp Exp (Aliases rep)
_ = forall a. Maybe a
Nothing

-- | This was former @createsAliasedArrOK@ from DataStructs
--   While Rearrange and Rotate create aliased arrays, we
--   do not yet support them because it would mean we have
--   to "reverse" the index function, for example to support
--   coalescing in the case below,
--       @let a = map f a0   @
--       @let b = transpose a@
--       @let y[4] = copy(b) @
--   we would need to assign to @a@ as index function, the
--   inverse of the transpose, such that, when creating @b@
--   by transposition we get a directly-mapped array, which
--   is expected by the copying in y[4].
--   For the moment we support only transposition and VName-expressions,
--     but rotations and full slices could also be supported.
--
-- This function complements 'getDirAliasFromExp' by returning a function that
-- applies the inverse index function transformation.
getInvAliasFromExp :: Exp (Aliases rep) -> InvAlias
getInvAliasFromExp :: forall rep. Exp (Aliases rep) -> InvAlias
getInvAliasFromExp (BasicOp (SubExp (Var VName
_))) = forall a. a -> Maybe a
Just forall a. a -> a
id
getInvAliasFromExp (BasicOp (Opaque OpaqueOp
_ (Var VName
_))) = forall a. a -> Maybe a
Just forall a. a -> a
id
getInvAliasFromExp (BasicOp Update {}) = forall a. a -> Maybe a
Just forall a. a -> a
id
getInvAliasFromExp (BasicOp (Rearrange [Int]
perm VName
_)) =
  let perm' :: [Int]
perm' = forall a. [Int] -> [a] -> [a]
IxFun.permuteInv [Int]
perm [Int
0 .. forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
perm forall a. Num a => a -> a -> a
- Int
1]
   in forall a. a -> Maybe a
Just (forall num. IntegralExp num => IxFun num -> [Int] -> IxFun num
`IxFun.permute` [Int]
perm')
getInvAliasFromExp Exp (Aliases rep)
_ = forall a. Maybe a
Nothing

class TopDownHelper inner where
  innerNonNegatives :: [VName] -> inner -> Names

  innerKnownLessThan :: inner -> [(VName, PrimExp VName)]

  scopeHelper :: inner -> Scope rep

instance TopDownHelper (SegOp lvl rep) where
  innerNonNegatives :: [VName] -> SegOp lvl rep -> Names
innerNonNegatives [VName]
_ SegOp lvl rep
op =
    forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (VName -> Names
oneName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace forall a b. (a -> b) -> a -> b
$ forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op

  innerKnownLessThan :: SegOp lvl rep -> [(VName, PrimExp VName)]
innerKnownLessThan SegOp lvl rep
op =
    forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int64) forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace forall a b. (a -> b) -> a -> b
$ forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op

  scopeHelper :: forall rep. SegOp lvl rep -> Scope rep
scopeHelper SegOp lvl rep
op = forall rep. SegSpace -> Scope rep
scopeOfSegSpace forall a b. (a -> b) -> a -> b
$ forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op

instance TopDownHelper (HostOp NoOp (Aliases GPUMem)) where
  innerNonNegatives :: [VName] -> HostOp NoOp (Aliases GPUMem) -> Names
innerNonNegatives [VName]
vs (SegOp SegOp SegLevel (Aliases GPUMem)
op) = forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs SegOp SegLevel (Aliases GPUMem)
op
  innerNonNegatives [VName
vname] (SizeOp (GetSize Name
_ SizeClass
_)) = VName -> Names
oneName VName
vname
  innerNonNegatives [VName
vname] (SizeOp (GetSizeMax SizeClass
_)) = VName -> Names
oneName VName
vname
  innerNonNegatives [VName]
_ HostOp NoOp (Aliases GPUMem)
_ = forall a. Monoid a => a
mempty

  innerKnownLessThan :: HostOp NoOp (Aliases GPUMem) -> [(VName, PrimExp VName)]
innerKnownLessThan (SegOp SegOp SegLevel (Aliases GPUMem)
op) = forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan SegOp SegLevel (Aliases GPUMem)
op
  innerKnownLessThan HostOp NoOp (Aliases GPUMem)
_ = forall a. Monoid a => a
mempty

  scopeHelper :: forall rep. HostOp NoOp (Aliases GPUMem) -> Scope rep
scopeHelper (SegOp SegOp SegLevel (Aliases GPUMem)
op) = forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper SegOp SegLevel (Aliases GPUMem)
op
  scopeHelper HostOp NoOp (Aliases GPUMem)
_ = forall a. Monoid a => a
mempty

instance TopDownHelper (inner (Aliases MCMem)) => TopDownHelper (MC.MCOp inner (Aliases MCMem)) where
  innerNonNegatives :: [VName] -> MCOp inner (Aliases MCMem) -> Names
innerNonNegatives [VName]
vs (ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
op) =
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs) Maybe (SegOp () (Aliases MCMem))
par_op
      forall a. Semigroup a => a -> a -> a
<> forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs SegOp () (Aliases MCMem)
op
  innerNonNegatives [VName]
vs (MC.OtherOp inner (Aliases MCMem)
op) =
    forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs inner (Aliases MCMem)
op
  innerKnownLessThan :: MCOp inner (Aliases MCMem) -> [(VName, PrimExp VName)]
innerKnownLessThan (ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
op) =
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan Maybe (SegOp () (Aliases MCMem))
par_op forall a. Semigroup a => a -> a -> a
<> forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan SegOp () (Aliases MCMem)
op
  innerKnownLessThan (MC.OtherOp inner (Aliases MCMem)
op) =
    forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan inner (Aliases MCMem)
op
  scopeHelper :: forall rep. MCOp inner (Aliases MCMem) -> Scope rep
scopeHelper (ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
op) =
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper Maybe (SegOp () (Aliases MCMem))
par_op forall a. Semigroup a => a -> a -> a
<> forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper SegOp () (Aliases MCMem)
op
  scopeHelper MC.OtherOp {} = forall a. Monoid a => a
mempty

instance TopDownHelper (NoOp rep) where
  innerNonNegatives :: [VName] -> NoOp rep -> Names
innerNonNegatives [VName]
_ NoOp rep
NoOp = forall a. Monoid a => a
mempty
  innerKnownLessThan :: NoOp rep -> [(VName, PrimExp VName)]
innerKnownLessThan NoOp rep
NoOp = forall a. Monoid a => a
mempty
  scopeHelper :: forall rep. NoOp rep -> Scope rep
scopeHelper NoOp rep
NoOp = forall a. Monoid a => a
mempty

-- | fills in the TopdownEnv table
updateTopdownEnv ::
  (ASTRep rep, Op rep ~ MemOp inner rep, TopDownHelper (inner (Aliases rep))) =>
  TopdownEnv rep ->
  Stm (Aliases rep) ->
  TopdownEnv rep
updateTopdownEnv :: forall rep (inner :: * -> *).
(ASTRep rep, Op rep ~ MemOp inner rep,
 TopDownHelper (inner (Aliases rep))) =>
TopdownEnv rep -> Stm (Aliases rep) -> TopdownEnv rep
updateTopdownEnv TopdownEnv rep
env stm :: Stm (Aliases rep)
stm@(Let (Pat [PatElem (LetDec (Aliases rep))
pe]) StmAux (ExpDec (Aliases rep))
_ (Op (Alloc (Var VName
vname) Space
sp))) =
  TopdownEnv rep
env
    { alloc :: AllocTab
alloc = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
pe) Space
sp forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> AllocTab
alloc TopdownEnv rep
env,
      scope :: ScopeTab rep
scope = forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm,
      nonNegatives :: Names
nonNegatives = forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> VName -> Names
oneName VName
vname
    }
updateTopdownEnv TopdownEnv rep
env stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ (Op (Inner inner (Aliases rep)
inner))) =
  TopdownEnv rep
env
    { scope :: ScopeTab rep
scope = forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm forall a. Semigroup a => a -> a -> a
<> forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper inner (Aliases rep)
inner,
      nonNegatives :: Names
nonNegatives = forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives (forall dec. Pat dec -> [VName]
patNames Pat (LetDec (Aliases rep))
pat) inner (Aliases rep)
inner,
      knownLessThan :: [(VName, PrimExp VName)]
knownLessThan = forall rep. TopdownEnv rep -> [(VName, PrimExp VName)]
knownLessThan TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan inner (Aliases rep)
inner
    }
updateTopdownEnv TopdownEnv rep
env (Let (Pat [PatElem (LetDec (Aliases rep))]
_) StmAux (ExpDec (Aliases rep))
_ (BasicOp (Assert SubExp
se ErrorMsg SubExp
_ (SrcLoc, [SrcLoc])
_))) =
  TopdownEnv rep
env {td_asserts :: [SubExp]
td_asserts = SubExp
se forall a. a -> [a] -> [a]
: forall rep. TopdownEnv rep -> [SubExp]
td_asserts TopdownEnv rep
env}
updateTopdownEnv TopdownEnv rep
env stm :: Stm (Aliases rep)
stm@(Let (Pat [PatElem (LetDec (Aliases rep))
pe]) StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e)
  | Just (VName
x, DirAlias
ixfn) <- forall rep. Exp (Aliases rep) -> Maybe (VName, DirAlias)
getDirAliasFromExp Exp (Aliases rep)
e =
      let ixfn_inv :: InvAlias
ixfn_inv = forall rep. Exp (Aliases rep) -> InvAlias
getInvAliasFromExp Exp (Aliases rep)
e
       in TopdownEnv rep
env
            { v_alias :: VarAliasTab
v_alias = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
pe) (VName
x, DirAlias
ixfn, InvAlias
ixfn_inv) (forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
env),
              scope :: ScopeTab rep
scope = forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm,
              nonNegatives :: Names
nonNegatives = forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> forall rep. Typed rep => Pat rep -> Names
nonNegativesInPat (forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm (Aliases rep)
stm)
            }
updateTopdownEnv TopdownEnv rep
env Stm (Aliases rep)
stm =
  TopdownEnv rep
env
    { scope :: ScopeTab rep
scope = forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm,
      nonNegatives :: Names
nonNegatives =
        forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env
          forall a. Semigroup a => a -> a -> a
<> forall rep. Typed rep => Pat rep -> Names
nonNegativesInPat (forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm (Aliases rep)
stm)
    }

nonNegativesInPat :: Typed rep => Pat rep -> Names
nonNegativesInPat :: forall rep. Typed rep => Pat rep -> Names
nonNegativesInPat (Pat [PatElem rep]
elems) =
  forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ([VName] -> Names
namesFromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe SubExp -> Maybe VName
subExpVar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> [SubExp]
arrayDims forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. Typed t => t -> Type
typeOf) [PatElem rep]
elems

-- | The topdown handler for loops.
updateTopdownEnvLoop :: TopdownEnv rep -> [(FParam rep, SubExp)] -> LoopForm (Aliases rep) -> TopdownEnv rep
updateTopdownEnvLoop :: forall rep.
TopdownEnv rep
-> [(FParam rep, SubExp)]
-> LoopForm (Aliases rep)
-> TopdownEnv rep
updateTopdownEnvLoop TopdownEnv rep
td_env [(FParam rep, SubExp)]
arginis LoopForm (Aliases rep)
lform =
  let scopetab :: ScopeTab rep
scopetab =
        forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env
          forall a. Semigroup a => a -> a -> a
<> forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(FParam rep, SubExp)]
arginis)
          forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf LoopForm (Aliases rep)
lform
      non_negatives :: Names
non_negatives =
        forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
td_env forall a. Semigroup a => a -> a -> a
<> case LoopForm (Aliases rep)
lform of
          ForLoop VName
v IntType
_ SubExp
_ [(LParam (Aliases rep), VName)]
_ -> VName -> Names
oneName VName
v
          LoopForm (Aliases rep)
_ -> forall a. Monoid a => a
mempty
      less_than :: [(VName, PrimExp VName)]
less_than =
        case LoopForm (Aliases rep)
lform of
          ForLoop VName
v IntType
_ SubExp
b [(LParam (Aliases rep), VName)]
_ -> [(VName
v, PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
Int64) SubExp
b)]
          LoopForm (Aliases rep)
_ -> forall a. Monoid a => a
mempty
   in TopdownEnv rep
td_env
        { scope :: ScopeTab rep
scope = ScopeTab rep
scopetab,
          nonNegatives :: Names
nonNegatives = Names
non_negatives,
          knownLessThan :: [(VName, PrimExp VName)]
knownLessThan = [(VName, PrimExp VName)]
less_than forall a. Semigroup a => a -> a -> a
<> forall rep. TopdownEnv rep -> [(VName, PrimExp VName)]
knownLessThan TopdownEnv rep
td_env
        }

-- | Get direct aliased index function.  Returns a triple of current memory
-- block to be coalesced, the destination memory block and the index function of
-- the access in the space of the destination block.
getDirAliasedIxfn :: HasMemBlock (Aliases rep) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn :: forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn TopdownEnv rep
td_env CoalsTab
coals_tab VName
x =
  case forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
x (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) of
    Just (MemBlock PrimType
_ Shape
_ VName
m_x IxFun
orig_ixfun) ->
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_x CoalsTab
coals_tab of
        Just CoalsEntry
coal_etry -> do
          (Coalesced CoalescedKind
_ (MemBlock PrimType
_ Shape
_ VName
m IxFun
ixf) FreeVarSubsts
_) <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab (forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
td_env) (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_etry) VName
x
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
m_x, VName
m, IxFun
ixf)
        Maybe CoalsEntry
Nothing ->
          -- This value is not subject to coalescing at the moment. Just return the
          -- original index function
          forall a. a -> Maybe a
Just (VName
m_x, VName
m_x, IxFun
orig_ixfun)
    Maybe ArrayMemBound
Nothing -> forall a. Maybe a
Nothing

-- | Like 'getDirAliasedIxfn', but this version returns 'Nothing' if the value
-- is not currently subject to coalescing.
getDirAliasedIxfn' :: HasMemBlock (Aliases rep) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn' :: forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn' TopdownEnv rep
td_env CoalsTab
coals_tab VName
x =
  case forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
x (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) of
    Just (MemBlock PrimType
_ Shape
_ VName
m_x IxFun
_) ->
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_x CoalsTab
coals_tab of
        Just CoalsEntry
coal_etry -> do
          (Coalesced CoalescedKind
_ (MemBlock PrimType
_ Shape
_ VName
m IxFun
ixf) FreeVarSubsts
_) <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab (forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
td_env) (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_etry) VName
x
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
m_x, VName
m, IxFun
ixf)
        Maybe CoalsEntry
Nothing ->
          -- This value is not subject to coalescing at the moment. Just return the
          -- original index function
          forall a. Maybe a
Nothing
    Maybe ArrayMemBound
Nothing -> forall a. Maybe a
Nothing

-- | Given a 'VName', walk the 'VarAliasTab' until found in the 'Map'.
walkAliasTab ::
  VarAliasTab ->
  M.Map VName Coalesced ->
  VName ->
  Maybe Coalesced
walkAliasTab :: VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab VarAliasTab
_ Map VName Coalesced
vtab VName
x
  | Just Coalesced
c <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Map VName Coalesced
vtab =
      forall a. a -> Maybe a
Just Coalesced
c -- @x@ is in @vartab@ together with its new ixfun
walkAliasTab VarAliasTab
alias_tab Map VName Coalesced
vtab VName
x
  | Just (VName
x0, DirAlias
alias0, InvAlias
_) <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x VarAliasTab
alias_tab = do
      Coalesced CoalescedKind
knd (MemBlock PrimType
pt Shape
shp VName
vname IxFun
ixf) FreeVarSubsts
substs <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab VarAliasTab
alias_tab Map VName Coalesced
vtab VName
x0
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
knd (PrimType -> Shape -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
pt Shape
shp VName
vname forall a b. (a -> b) -> a -> b
$ DirAlias
alias0 IxFun
ixf) FreeVarSubsts
substs
walkAliasTab VarAliasTab
_ Map VName Coalesced
_ VName
_ = forall a. Maybe a
Nothing

-- | We assume @x@ is in @vartab@ and we add the variables that @x@ aliases
--   for as long as possible following a chain of direct-aliasing operators,
--   i.e., without considering aliasing of if-then-else, loops, etc. For example:
--     @ x0 = if c then ... else ...@
--     @ x1 = rearrange r1 x0 @
--     @ x2 = reverse x1@
--     @ y[slc] = x2 @
--   We assume @vartab@ constains a binding for @x2@, and calling this function
--     with @x2@ as argument should also insert entries for @x1@ and @x0@ to
--     @vartab@, of course if their aliasing operations are invertible.
--   We assume inverting aliases has been performed by the top-down pass.
addInvAliassesVarTab ::
  HasMemBlock (Aliases rep) =>
  TopdownEnv rep ->
  M.Map VName Coalesced ->
  VName ->
  Maybe (M.Map VName Coalesced)
addInvAliassesVarTab :: forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliassesVarTab TopdownEnv rep
td_env Map VName Coalesced
vtab VName
x
  | Just (Coalesced CoalescedKind
_ (MemBlock PrimType
_ Shape
_ VName
m_y IxFun
x_ixfun) FreeVarSubsts
fv_subs) <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Map VName Coalesced
vtab =
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x (forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
td_env) of
        Maybe (VName, DirAlias, InvAlias)
Nothing -> forall a. a -> Maybe a
Just Map VName Coalesced
vtab
        Just (VName
_, DirAlias
_, InvAlias
Nothing) -> forall a. Maybe a
Nothing -- can't invert ixfun, conservatively fail!
        Just (VName
x0, DirAlias
_, Just DirAlias
inv_alias0) ->
          let x_ixfn0 :: IxFun
x_ixfn0 = DirAlias
inv_alias0 IxFun
x_ixfun
           in case forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
x0 (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) of
                Maybe ArrayMemBound
Nothing -> forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"
                Just (MemBlock PrimType
ptp Shape
shp VName
_ IxFun
_) ->
                  let coal :: Coalesced
coal = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
TransitiveCoal (PrimType -> Shape -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
ptp Shape
shp VName
m_y IxFun
x_ixfn0) FreeVarSubsts
fv_subs
                      vartab' :: Map VName Coalesced
vartab' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
x0 Coalesced
coal Map VName Coalesced
vtab
                   in forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliassesVarTab TopdownEnv rep
td_env Map VName Coalesced
vartab' VName
x0
addInvAliassesVarTab TopdownEnv rep
_ Map VName Coalesced
_ VName
_ = forall a. Maybe a
Nothing

areAliased :: TopdownEnv rep -> VName -> VName -> Bool
areAliased :: forall rep. TopdownEnv rep -> VName -> VName -> Bool
areAliased TopdownEnv rep
_ VName
m_x VName
m_y =
  -- this is a dummy implementation
  VName
m_x forall a. Eq a => a -> a -> Bool
== VName
m_y

areAnyAliased :: TopdownEnv rep -> VName -> [VName] -> Bool
areAnyAliased :: forall rep. TopdownEnv rep -> VName -> [VName] -> Bool
areAnyAliased TopdownEnv rep
td_env VName
m_x =
  forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall rep. TopdownEnv rep -> VName -> VName -> Bool
areAliased TopdownEnv rep
td_env VName
m_x)