{-# LANGUAGE LambdaCase #-}

module Futhark.Analysis.AccessPattern
  ( analyseDimAccesses,
    analyseFunction,
    vnameFromSegOp,
    analysisPropagateByTransitivity,
    isInvariant,
    Analyse,
    IndexTable,
    ArrayName,
    DimAccess (..),
    IndexExprName,
    BodyType (..),
    SegOpName (SegmentedMap, SegmentedRed, SegmentedScan, SegmentedHist),
    Context (..),
    analyseIndex,
    VariableInfo (..),
    VarType (..),
    isCounter,
    Dependency (..),
  )
where

import Data.Bifunctor
import Data.Foldable
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.IR.Aliases
import Futhark.IR.GPU
import Futhark.IR.GPUMem
import Futhark.IR.MC
import Futhark.IR.MCMem
import Futhark.IR.SOACS
import Futhark.IR.Seq
import Futhark.IR.SeqMem
import Futhark.Util.Pretty

-- | Name of a SegOp, used to identify the SegOp that an array access is
-- contained in.
data SegOpName
  = SegmentedMap {SegOpName -> VName
vnameFromSegOp :: VName}
  | SegmentedRed {vnameFromSegOp :: VName}
  | SegmentedScan {vnameFromSegOp :: VName}
  | SegmentedHist {vnameFromSegOp :: VName}
  deriving (SegOpName -> SegOpName -> Bool
(SegOpName -> SegOpName -> Bool)
-> (SegOpName -> SegOpName -> Bool) -> Eq SegOpName
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SegOpName -> SegOpName -> Bool
== :: SegOpName -> SegOpName -> Bool
$c/= :: SegOpName -> SegOpName -> Bool
/= :: SegOpName -> SegOpName -> Bool
Eq, Eq SegOpName
Eq SegOpName =>
(SegOpName -> SegOpName -> Ordering)
-> (SegOpName -> SegOpName -> Bool)
-> (SegOpName -> SegOpName -> Bool)
-> (SegOpName -> SegOpName -> Bool)
-> (SegOpName -> SegOpName -> Bool)
-> (SegOpName -> SegOpName -> SegOpName)
-> (SegOpName -> SegOpName -> SegOpName)
-> Ord SegOpName
SegOpName -> SegOpName -> Bool
SegOpName -> SegOpName -> Ordering
SegOpName -> SegOpName -> SegOpName
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: SegOpName -> SegOpName -> Ordering
compare :: SegOpName -> SegOpName -> Ordering
$c< :: SegOpName -> SegOpName -> Bool
< :: SegOpName -> SegOpName -> Bool
$c<= :: SegOpName -> SegOpName -> Bool
<= :: SegOpName -> SegOpName -> Bool
$c> :: SegOpName -> SegOpName -> Bool
> :: SegOpName -> SegOpName -> Bool
$c>= :: SegOpName -> SegOpName -> Bool
>= :: SegOpName -> SegOpName -> Bool
$cmax :: SegOpName -> SegOpName -> SegOpName
max :: SegOpName -> SegOpName -> SegOpName
$cmin :: SegOpName -> SegOpName -> SegOpName
min :: SegOpName -> SegOpName -> SegOpName
Ord, Int -> SegOpName -> ShowS
[SegOpName] -> ShowS
SegOpName -> [Char]
(Int -> SegOpName -> ShowS)
-> (SegOpName -> [Char])
-> ([SegOpName] -> ShowS)
-> Show SegOpName
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SegOpName -> ShowS
showsPrec :: Int -> SegOpName -> ShowS
$cshow :: SegOpName -> [Char]
show :: SegOpName -> [Char]
$cshowList :: [SegOpName] -> ShowS
showList :: [SegOpName] -> ShowS
Show)

-- | Name of an array indexing expression. Taken from the pattern of
-- the expression.
type IndexExprName = VName

data BodyType
  = SegOpName SegOpName
  | LoopBodyName VName
  | CondBodyName VName
  deriving (Int -> BodyType -> ShowS
[BodyType] -> ShowS
BodyType -> [Char]
(Int -> BodyType -> ShowS)
-> (BodyType -> [Char]) -> ([BodyType] -> ShowS) -> Show BodyType
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BodyType -> ShowS
showsPrec :: Int -> BodyType -> ShowS
$cshow :: BodyType -> [Char]
show :: BodyType -> [Char]
$cshowList :: [BodyType] -> ShowS
showList :: [BodyType] -> ShowS
Show, Eq BodyType
Eq BodyType =>
(BodyType -> BodyType -> Ordering)
-> (BodyType -> BodyType -> Bool)
-> (BodyType -> BodyType -> Bool)
-> (BodyType -> BodyType -> Bool)
-> (BodyType -> BodyType -> Bool)
-> (BodyType -> BodyType -> BodyType)
-> (BodyType -> BodyType -> BodyType)
-> Ord BodyType
BodyType -> BodyType -> Bool
BodyType -> BodyType -> Ordering
BodyType -> BodyType -> BodyType
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: BodyType -> BodyType -> Ordering
compare :: BodyType -> BodyType -> Ordering
$c< :: BodyType -> BodyType -> Bool
< :: BodyType -> BodyType -> Bool
$c<= :: BodyType -> BodyType -> Bool
<= :: BodyType -> BodyType -> Bool
$c> :: BodyType -> BodyType -> Bool
> :: BodyType -> BodyType -> Bool
$c>= :: BodyType -> BodyType -> Bool
>= :: BodyType -> BodyType -> Bool
$cmax :: BodyType -> BodyType -> BodyType
max :: BodyType -> BodyType -> BodyType
$cmin :: BodyType -> BodyType -> BodyType
min :: BodyType -> BodyType -> BodyType
Ord, BodyType -> BodyType -> Bool
(BodyType -> BodyType -> Bool)
-> (BodyType -> BodyType -> Bool) -> Eq BodyType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BodyType -> BodyType -> Bool
== :: BodyType -> BodyType -> Bool
$c/= :: BodyType -> BodyType -> Bool
/= :: BodyType -> BodyType -> Bool
Eq)

-- | Stores the name of an array, the nest of loops, kernels,
-- conditionals in which it is constructed, and the existing layout of
-- the array. The latter is currently largely unused and not
-- trustworthy, but might be useful in the future.
type ArrayName = (VName, [BodyType], [Int])

-- | Tuple of patternName and nested `level` it index occurred at, as well as
-- what the actual iteration type is.
data Dependency = Dependency
  { Dependency -> Int
lvl :: Int,
    Dependency -> VarType
varType :: VarType
  }
  deriving (Dependency -> Dependency -> Bool
(Dependency -> Dependency -> Bool)
-> (Dependency -> Dependency -> Bool) -> Eq Dependency
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Dependency -> Dependency -> Bool
== :: Dependency -> Dependency -> Bool
$c/= :: Dependency -> Dependency -> Bool
/= :: Dependency -> Dependency -> Bool
Eq, Int -> Dependency -> ShowS
[Dependency] -> ShowS
Dependency -> [Char]
(Int -> Dependency -> ShowS)
-> (Dependency -> [Char])
-> ([Dependency] -> ShowS)
-> Show Dependency
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Dependency -> ShowS
showsPrec :: Int -> Dependency -> ShowS
$cshow :: Dependency -> [Char]
show :: Dependency -> [Char]
$cshowList :: [Dependency] -> ShowS
showList :: [Dependency] -> ShowS
Show)

-- | Collect all features of access to a specific dimension of an array.
data DimAccess rep = DimAccess
  { -- | Set of VNames of iteration variables (gtids, loop counters, etc.)
    -- that some access is variant to.
    -- An empty set indicates that the access is invariant.
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies :: M.Map VName Dependency,
    -- | Used to store the name of the original expression from which `dependencies`
    -- was computed. `Nothing` if it is a constant.
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
originalVar :: Maybe VName
  }
  deriving (DimAccess rep -> DimAccess rep -> Bool
(DimAccess rep -> DimAccess rep -> Bool)
-> (DimAccess rep -> DimAccess rep -> Bool) -> Eq (DimAccess rep)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (rep :: k). DimAccess rep -> DimAccess rep -> Bool
$c== :: forall k (rep :: k). DimAccess rep -> DimAccess rep -> Bool
== :: DimAccess rep -> DimAccess rep -> Bool
$c/= :: forall k (rep :: k). DimAccess rep -> DimAccess rep -> Bool
/= :: DimAccess rep -> DimAccess rep -> Bool
Eq, Int -> DimAccess rep -> ShowS
[DimAccess rep] -> ShowS
DimAccess rep -> [Char]
(Int -> DimAccess rep -> ShowS)
-> (DimAccess rep -> [Char])
-> ([DimAccess rep] -> ShowS)
-> Show (DimAccess rep)
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
forall k (rep :: k). Int -> DimAccess rep -> ShowS
forall k (rep :: k). [DimAccess rep] -> ShowS
forall k (rep :: k). DimAccess rep -> [Char]
$cshowsPrec :: forall k (rep :: k). Int -> DimAccess rep -> ShowS
showsPrec :: Int -> DimAccess rep -> ShowS
$cshow :: forall k (rep :: k). DimAccess rep -> [Char]
show :: DimAccess rep -> [Char]
$cshowList :: forall k (rep :: k). [DimAccess rep] -> ShowS
showList :: [DimAccess rep] -> ShowS
Show)

instance Semigroup (DimAccess rep) where
  DimAccess rep
adeps <> :: DimAccess rep -> DimAccess rep -> DimAccess rep
<> DimAccess rep
bdeps =
    Map VName Dependency -> Maybe VName -> DimAccess rep
forall {k} (rep :: k).
Map VName Dependency -> Maybe VName -> DimAccess rep
DimAccess
      (DimAccess rep -> Map VName Dependency
forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies DimAccess rep
adeps Map VName Dependency
-> Map VName Dependency -> Map VName Dependency
forall a. Semigroup a => a -> a -> a
<> DimAccess rep -> Map VName Dependency
forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies DimAccess rep
bdeps)
      ( case DimAccess rep -> Maybe VName
forall {k} (rep :: k). DimAccess rep -> Maybe VName
originalVar DimAccess rep
adeps of
          Maybe VName
Nothing -> DimAccess rep -> Maybe VName
forall {k} (rep :: k). DimAccess rep -> Maybe VName
originalVar DimAccess rep
bdeps
          Maybe VName
_ -> DimAccess rep -> Maybe VName
forall {k} (rep :: k). DimAccess rep -> Maybe VName
originalVar DimAccess rep
adeps
      )

instance Monoid (DimAccess rep) where
  mempty :: DimAccess rep
mempty = Map VName Dependency -> Maybe VName -> DimAccess rep
forall {k} (rep :: k).
Map VName Dependency -> Maybe VName -> DimAccess rep
DimAccess Map VName Dependency
forall a. Monoid a => a
mempty Maybe VName
forall a. Maybe a
Nothing

isInvariant :: DimAccess rep -> Bool
isInvariant :: forall {k} (rep :: k). DimAccess rep -> Bool
isInvariant = Map VName Dependency -> Bool
forall a. Map VName a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Map VName Dependency -> Bool)
-> (DimAccess rep -> Map VName Dependency) -> DimAccess rep -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DimAccess rep -> Map VName Dependency
forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies

-- | For each array access in a program, this data structure stores the
-- dependencies of each dimension in the access, the array name, and the
-- name of the SegOp that the access is contained in.
-- Each DimAccess element corresponds to an access to a given dimension
-- in the given array, in the same order of the dimensions.
type IndexTable rep =
  M.Map SegOpName (M.Map ArrayName (M.Map IndexExprName [DimAccess rep]))

unionIndexTables :: IndexTable rep -> IndexTable rep -> IndexTable rep
unionIndexTables :: forall {k} (rep :: k).
IndexTable rep -> IndexTable rep -> IndexTable rep
unionIndexTables = (Map ArrayName (Map VName [DimAccess rep])
 -> Map ArrayName (Map VName [DimAccess rep])
 -> Map ArrayName (Map VName [DimAccess rep]))
-> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
-> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
-> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith ((Map VName [DimAccess rep]
 -> Map VName [DimAccess rep] -> Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Map VName [DimAccess rep]
-> Map VName [DimAccess rep] -> Map VName [DimAccess rep]
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union)

-- | Make segops on arrays transitive, ie. if
-- > let A = segmap (..) xs -- A indexes into xs
-- > let B = segmap (..) A  -- B indexes into A
-- Then B also derives all A's array-accesses, like xs.
-- Runs in n²
analysisPropagateByTransitivity :: IndexTable rep -> IndexTable rep
analysisPropagateByTransitivity :: forall {k} (rep :: k). IndexTable rep -> IndexTable rep
analysisPropagateByTransitivity IndexTable rep
idx_table =
  (Map ArrayName (Map VName [DimAccess rep])
 -> Map ArrayName (Map VName [DimAccess rep]))
-> IndexTable rep -> IndexTable rep
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Map ArrayName (Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
foldlArrayNameMap IndexTable rep
idx_table
  where
    aggregateResults :: VName -> Map ArrayName (Map VName [DimAccess rep])
aggregateResults VName
arr_name =
      Map ArrayName (Map VName [DimAccess rep])
-> (Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep]))
-> Maybe (Map ArrayName (Map VName [DimAccess rep]))
-> Map ArrayName (Map VName [DimAccess rep])
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
        Map ArrayName (Map VName [DimAccess rep])
forall a. Monoid a => a
mempty
        Map ArrayName (Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
foldlArrayNameMap
        ((SegOpName -> VName)
-> IndexTable rep
-> Map VName (Map ArrayName (Map VName [DimAccess rep]))
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys SegOpName -> VName
vnameFromSegOp IndexTable rep
idx_table Map VName (Map ArrayName (Map VName [DimAccess rep]))
-> VName -> Maybe (Map ArrayName (Map VName [DimAccess rep]))
forall k a. Ord k => Map k a -> k -> Maybe a
M.!? VName
arr_name)

    foldlArrayNameMap :: Map ArrayName (Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
foldlArrayNameMap Map ArrayName (Map VName [DimAccess rep])
aMap =
      (Map ArrayName (Map VName [DimAccess rep])
 -> Map ArrayName (Map VName [DimAccess rep])
 -> Map ArrayName (Map VName [DimAccess rep]))
-> Map ArrayName (Map VName [DimAccess rep])
-> [Map ArrayName (Map VName [DimAccess rep])]
-> Map ArrayName (Map VName [DimAccess rep])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((Map VName [DimAccess rep]
 -> Map VName [DimAccess rep] -> Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
-> Map ArrayName (Map VName [DimAccess rep])
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Map VName [DimAccess rep]
-> Map VName [DimAccess rep] -> Map VName [DimAccess rep]
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union) Map ArrayName (Map VName [DimAccess rep])
aMap ([Map ArrayName (Map VName [DimAccess rep])]
 -> Map ArrayName (Map VName [DimAccess rep]))
-> [Map ArrayName (Map VName [DimAccess rep])]
-> Map ArrayName (Map VName [DimAccess rep])
forall a b. (a -> b) -> a -> b
$
        (ArrayName -> Map ArrayName (Map VName [DimAccess rep]))
-> [ArrayName] -> [Map ArrayName (Map VName [DimAccess rep])]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Map ArrayName (Map VName [DimAccess rep])
aggregateResults (VName -> Map ArrayName (Map VName [DimAccess rep]))
-> (ArrayName -> VName)
-> ArrayName
-> Map ArrayName (Map VName [DimAccess rep])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \(VName
a, [BodyType]
_, [Int]
_) -> VName
a) ([ArrayName] -> [Map ArrayName (Map VName [DimAccess rep])])
-> [ArrayName] -> [Map ArrayName (Map VName [DimAccess rep])]
forall a b. (a -> b) -> a -> b
$
          Map ArrayName (Map VName [DimAccess rep]) -> [ArrayName]
forall k a. Map k a -> [k]
M.keys Map ArrayName (Map VName [DimAccess rep])
aMap

--
-- Helper types and functions to perform the analysis.
--

-- | Used during the analysis to keep track of the dependencies of patterns
-- encountered so far.
data Context rep = Context
  { -- | A mapping from patterns occuring in Let expressions to their dependencies
    --  and iteration types.
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments :: M.Map VName (VariableInfo rep),
    -- | Maps from sliced arrays to their respective access patterns.
    forall {k} (rep :: k).
Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
slices :: M.Map IndexExprName (ArrayName, [VName], [DimAccess rep]),
    -- | A list of the segMaps encountered during the analysis in the order they
    -- were encountered.
    forall {k} (rep :: k). Context rep -> [BodyType]
parents :: [BodyType],
    -- | Current level of recursion, also just `length parents`
    forall {k} (rep :: k). Context rep -> Int
currentLevel :: Int
  }
  deriving (Int -> Context rep -> ShowS
[Context rep] -> ShowS
Context rep -> [Char]
(Int -> Context rep -> ShowS)
-> (Context rep -> [Char])
-> ([Context rep] -> ShowS)
-> Show (Context rep)
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
forall k (rep :: k). Int -> Context rep -> ShowS
forall k (rep :: k). [Context rep] -> ShowS
forall k (rep :: k). Context rep -> [Char]
$cshowsPrec :: forall k (rep :: k). Int -> Context rep -> ShowS
showsPrec :: Int -> Context rep -> ShowS
$cshow :: forall k (rep :: k). Context rep -> [Char]
show :: Context rep -> [Char]
$cshowList :: forall k (rep :: k). [Context rep] -> ShowS
showList :: [Context rep] -> ShowS
Show, Context rep -> Context rep -> Bool
(Context rep -> Context rep -> Bool)
-> (Context rep -> Context rep -> Bool) -> Eq (Context rep)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (rep :: k). Context rep -> Context rep -> Bool
$c== :: forall k (rep :: k). Context rep -> Context rep -> Bool
== :: Context rep -> Context rep -> Bool
$c/= :: forall k (rep :: k). Context rep -> Context rep -> Bool
/= :: Context rep -> Context rep -> Bool
Eq)

instance Monoid (Context rep) where
  mempty :: Context rep
mempty =
    Context
      { assignments :: Map VName (VariableInfo rep)
assignments = Map VName (VariableInfo rep)
forall a. Monoid a => a
mempty,
        slices :: Map VName (ArrayName, [VName], [DimAccess rep])
slices = Map VName (ArrayName, [VName], [DimAccess rep])
forall a. Monoid a => a
mempty,
        parents :: [BodyType]
parents = [],
        currentLevel :: Int
currentLevel = Int
0
      }

instance Semigroup (Context rep) where
  Context Map VName (VariableInfo rep)
ass0 Map VName (ArrayName, [VName], [DimAccess rep])
slices0 [BodyType]
lastBody0 Int
lvl0 <> :: Context rep -> Context rep -> Context rep
<> Context Map VName (VariableInfo rep)
ass1 Map VName (ArrayName, [VName], [DimAccess rep])
slices1 [BodyType]
lastBody1 Int
lvl1 =
    Map VName (VariableInfo rep)
-> Map VName (ArrayName, [VName], [DimAccess rep])
-> [BodyType]
-> Int
-> Context rep
forall {k} (rep :: k).
Map VName (VariableInfo rep)
-> Map VName (ArrayName, [VName], [DimAccess rep])
-> [BodyType]
-> Int
-> Context rep
Context
      (Map VName (VariableInfo rep)
ass0 Map VName (VariableInfo rep)
-> Map VName (VariableInfo rep) -> Map VName (VariableInfo rep)
forall a. Semigroup a => a -> a -> a
<> Map VName (VariableInfo rep)
ass1)
      (Map VName (ArrayName, [VName], [DimAccess rep])
slices0 Map VName (ArrayName, [VName], [DimAccess rep])
-> Map VName (ArrayName, [VName], [DimAccess rep])
-> Map VName (ArrayName, [VName], [DimAccess rep])
forall a. Semigroup a => a -> a -> a
<> Map VName (ArrayName, [VName], [DimAccess rep])
slices1)
      ([BodyType]
lastBody0 [BodyType] -> [BodyType] -> [BodyType]
forall a. Semigroup a => a -> a -> a
<> [BodyType]
lastBody1)
      (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
lvl0 Int
lvl1)

-- | Extend a context with another context.
-- We never have to consider the case where VNames clash in the context, since
-- they are unique.
extend :: Context rep -> Context rep -> Context rep
extend :: forall k (rep :: k). Context rep -> Context rep -> Context rep
extend = Context rep -> Context rep -> Context rep
forall a. Semigroup a => a -> a -> a
(<>)

allSegMap :: Context rep -> [SegOpName]
allSegMap :: forall {k} (rep :: k). Context rep -> [SegOpName]
allSegMap (Context Map VName (VariableInfo rep)
_ Map VName (ArrayName, [VName], [DimAccess rep])
_ [BodyType]
parents Int
_) = (BodyType -> Maybe SegOpName) -> [BodyType] -> [SegOpName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe BodyType -> Maybe SegOpName
f [BodyType]
parents
  where
    f :: BodyType -> Maybe SegOpName
f (SegOpName SegOpName
o) = SegOpName -> Maybe SegOpName
forall a. a -> Maybe a
Just SegOpName
o
    f BodyType
_ = Maybe SegOpName
forall a. Maybe a
Nothing

-- | Context Value (VariableInfo) is the type used in the context to categorize
-- assignments. For example, a pattern might depend on a function parameter, a
-- gtid, or some other pattern.
data VariableInfo rep = VariableInfo
  { forall {k} (rep :: k). VariableInfo rep -> Names
deps :: Names,
    forall {k} (rep :: k). VariableInfo rep -> Int
level :: Int,
    forall {k} (rep :: k). VariableInfo rep -> [BodyType]
parents_nest :: [BodyType],
    forall {k} (rep :: k). VariableInfo rep -> VarType
variableType :: VarType
  }
  deriving (Int -> VariableInfo rep -> ShowS
[VariableInfo rep] -> ShowS
VariableInfo rep -> [Char]
(Int -> VariableInfo rep -> ShowS)
-> (VariableInfo rep -> [Char])
-> ([VariableInfo rep] -> ShowS)
-> Show (VariableInfo rep)
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
forall k (rep :: k). Int -> VariableInfo rep -> ShowS
forall k (rep :: k). [VariableInfo rep] -> ShowS
forall k (rep :: k). VariableInfo rep -> [Char]
$cshowsPrec :: forall k (rep :: k). Int -> VariableInfo rep -> ShowS
showsPrec :: Int -> VariableInfo rep -> ShowS
$cshow :: forall k (rep :: k). VariableInfo rep -> [Char]
show :: VariableInfo rep -> [Char]
$cshowList :: forall k (rep :: k). [VariableInfo rep] -> ShowS
showList :: [VariableInfo rep] -> ShowS
Show, VariableInfo rep -> VariableInfo rep -> Bool
(VariableInfo rep -> VariableInfo rep -> Bool)
-> (VariableInfo rep -> VariableInfo rep -> Bool)
-> Eq (VariableInfo rep)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (rep :: k). VariableInfo rep -> VariableInfo rep -> Bool
$c== :: forall k (rep :: k). VariableInfo rep -> VariableInfo rep -> Bool
== :: VariableInfo rep -> VariableInfo rep -> Bool
$c/= :: forall k (rep :: k). VariableInfo rep -> VariableInfo rep -> Bool
/= :: VariableInfo rep -> VariableInfo rep -> Bool
Eq)

data VarType
  = ConstType
  | Variable
  | ThreadID
  | LoopVar
  deriving (Int -> VarType -> ShowS
[VarType] -> ShowS
VarType -> [Char]
(Int -> VarType -> ShowS)
-> (VarType -> [Char]) -> ([VarType] -> ShowS) -> Show VarType
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> VarType -> ShowS
showsPrec :: Int -> VarType -> ShowS
$cshow :: VarType -> [Char]
show :: VarType -> [Char]
$cshowList :: [VarType] -> ShowS
showList :: [VarType] -> ShowS
Show, VarType -> VarType -> Bool
(VarType -> VarType -> Bool)
-> (VarType -> VarType -> Bool) -> Eq VarType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: VarType -> VarType -> Bool
== :: VarType -> VarType -> Bool
$c/= :: VarType -> VarType -> Bool
/= :: VarType -> VarType -> Bool
Eq)

isCounter :: VarType -> Bool
isCounter :: VarType -> Bool
isCounter VarType
LoopVar = Bool
True
isCounter VarType
ThreadID = Bool
True
isCounter VarType
_ = Bool
False

varInfoFromNames :: Context rep -> Names -> VariableInfo rep
varInfoFromNames :: forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx Names
names = do
  Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
forall {k} (rep :: k).
Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
VariableInfo Names
names (Context rep -> Int
forall {k} (rep :: k). Context rep -> Int
currentLevel Context rep
ctx) (Context rep -> [BodyType]
forall {k} (rep :: k). Context rep -> [BodyType]
parents Context rep
ctx) VarType
Variable

-- | Wrapper around the constructur of Context.
oneContext :: VName -> VariableInfo rep -> Context rep
oneContext :: forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
oneContext VName
name VariableInfo rep
var_info =
  Context
    { assignments :: Map VName (VariableInfo rep)
assignments = VName -> VariableInfo rep -> Map VName (VariableInfo rep)
forall k a. k -> a -> Map k a
M.singleton VName
name VariableInfo rep
var_info,
      slices :: Map VName (ArrayName, [VName], [DimAccess rep])
slices = Map VName (ArrayName, [VName], [DimAccess rep])
forall a. Monoid a => a
mempty,
      parents :: [BodyType]
parents = [],
      currentLevel :: Int
currentLevel = Int
0
    }

-- | Create a singular varInfo with no dependencies.
varInfoZeroDeps :: Context rep -> VariableInfo rep
varInfoZeroDeps :: forall {k} (rep :: k). Context rep -> VariableInfo rep
varInfoZeroDeps Context rep
ctx =
  Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
forall {k} (rep :: k).
Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
VariableInfo Names
forall a. Monoid a => a
mempty (Context rep -> Int
forall {k} (rep :: k). Context rep -> Int
currentLevel Context rep
ctx) (Context rep -> [BodyType]
forall {k} (rep :: k). Context rep -> [BodyType]
parents Context rep
ctx) VarType
Variable

-- | Create a singular context from a segspace
contextFromNames :: Context rep -> VariableInfo rep -> [VName] -> Context rep
contextFromNames :: forall {k} (rep :: k).
Context rep -> VariableInfo rep -> [VName] -> Context rep
contextFromNames Context rep
ctx VariableInfo rep
var_info = (Context rep -> Context rep -> Context rep)
-> Context rep -> [Context rep] -> Context rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Context rep -> Context rep -> Context rep
forall k (rep :: k). Context rep -> Context rep -> Context rep
extend Context rep
ctx ([Context rep] -> Context rep)
-> ([VName] -> [Context rep]) -> [VName] -> Context rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Context rep) -> [VName] -> [Context rep]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> VariableInfo rep -> Context rep
forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
`oneContext` VariableInfo rep
var_info)

-- | A representation where we can analyse access patterns.
class Analyse rep where
  -- | Analyse the op for this representation.
  analyseOp :: Op rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)

-- | Analyse each `entry` and accumulate the results.
analyseDimAccesses :: (Analyse rep) => Prog rep -> IndexTable rep
analyseDimAccesses :: forall rep. Analyse rep => Prog rep -> IndexTable rep
analyseDimAccesses = (FunDef rep -> IndexTable rep) -> [FunDef rep] -> IndexTable rep
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap' FunDef rep -> IndexTable rep
forall rep. Analyse rep => FunDef rep -> IndexTable rep
analyseFunction ([FunDef rep] -> IndexTable rep)
-> (Prog rep -> [FunDef rep]) -> Prog rep -> IndexTable rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog rep -> [FunDef rep]
forall rep. Prog rep -> [FunDef rep]
progFuns

-- | Analyse each statement in a function body.
analyseFunction :: (Analyse rep) => FunDef rep -> IndexTable rep
analyseFunction :: forall rep. Analyse rep => FunDef rep -> IndexTable rep
analyseFunction FunDef rep
func =
  let stms :: [Stm rep]
stms = Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms rep -> [Stm rep])
-> (Body rep -> Stms rep) -> Body rep -> [Stm rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms (Body rep -> [Stm rep]) -> Body rep -> [Stm rep]
forall a b. (a -> b) -> a -> b
$ FunDef rep -> Body rep
forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
func
      -- Create a context containing the function parameters
      ctx :: Context rep
ctx = Context rep -> VariableInfo rep -> [VName] -> Context rep
forall {k} (rep :: k).
Context rep -> VariableInfo rep -> [VName] -> Context rep
contextFromNames Context rep
forall a. Monoid a => a
mempty (Context rep -> VariableInfo rep
forall {k} (rep :: k). Context rep -> VariableInfo rep
varInfoZeroDeps Context rep
ctx) ([VName] -> Context rep) -> [VName] -> Context rep
forall a b. (a -> b) -> a -> b
$ (Param (FParamInfo rep) -> VName)
-> [Param (FParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName ([Param (FParamInfo rep)] -> [VName])
-> [Param (FParamInfo rep)] -> [VName]
forall a b. (a -> b) -> a -> b
$ FunDef rep -> [Param (FParamInfo rep)]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef rep
func
   in (Context rep, IndexTable rep) -> IndexTable rep
forall a b. (a, b) -> b
snd ((Context rep, IndexTable rep) -> IndexTable rep)
-> (Context rep, IndexTable rep) -> IndexTable rep
forall a b. (a -> b) -> a -> b
$ Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
analyseStmsPrimitive Context rep
forall k (rep :: k). Context rep
ctx [Stm rep]
stms

-- | Analyse each statement in a list of statements.
analyseStmsPrimitive :: (Analyse rep) => Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
analyseStmsPrimitive :: forall rep.
Analyse rep =>
Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
analyseStmsPrimitive Context rep
ctx =
  -- Fold over statements in body
  ((Context rep, IndexTable rep)
 -> Stm rep -> (Context rep, IndexTable rep))
-> (Context rep, IndexTable rep)
-> [Stm rep]
-> (Context rep, IndexTable rep)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
    (\(Context rep
c, IndexTable rep
r) Stm rep
stm -> (IndexTable rep -> IndexTable rep)
-> (Context rep, IndexTable rep) -> (Context rep, IndexTable rep)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (IndexTable rep -> IndexTable rep -> IndexTable rep
forall {k} (rep :: k).
IndexTable rep -> IndexTable rep -> IndexTable rep
unionIndexTables IndexTable rep
r) ((Context rep, IndexTable rep) -> (Context rep, IndexTable rep))
-> (Context rep, IndexTable rep) -> (Context rep, IndexTable rep)
forall a b. (a -> b) -> a -> b
$ Context rep -> Stm rep -> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep -> Stm rep -> (Context rep, IndexTable rep)
analyseStm Context rep
c Stm rep
stm)
    (Context rep
ctx, IndexTable rep
forall a. Monoid a => a
mempty)

-- | Same as analyseStmsPrimitive, but change the resulting context into
-- a varInfo, mapped to pattern.
analyseStms :: (Analyse rep) => Context rep -> (VName -> BodyType) -> [VName] -> [Stm rep] -> (Context rep, IndexTable rep)
analyseStms :: forall rep.
Analyse rep =>
Context rep
-> (VName -> BodyType)
-> [VName]
-> [Stm rep]
-> (Context rep, IndexTable rep)
analyseStms Context rep
ctx VName -> BodyType
body_constructor [VName]
pats [Stm rep]
body = do
  -- 0. Recurse into body with ctx
  let (Context rep
ctx'', IndexTable rep
indexTable) = Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
analyseStmsPrimitive Context rep
recContext [Stm rep]
body

  -- 0.1 Get all new slices
  let slices_new :: Map VName (ArrayName, [VName], [DimAccess rep])
slices_new = Map VName (ArrayName, [VName], [DimAccess rep])
-> Map VName (ArrayName, [VName], [DimAccess rep])
-> Map VName (ArrayName, [VName], [DimAccess rep])
forall k a b. Ord k => Map k a -> Map k b -> Map k a
M.difference (Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
forall {k} (rep :: k).
Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
slices Context rep
ctx'') (Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
forall {k} (rep :: k).
Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
slices Context rep
ctx)
  -- 0.2 Make "IndexExpressions" of the slices
  let slices_indices :: IndexTable rep
slices_indices =
        (IndexTable rep -> IndexTable rep -> IndexTable rep)
-> IndexTable rep -> [IndexTable rep] -> IndexTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl IndexTable rep -> IndexTable rep -> IndexTable rep
forall {k} (rep :: k).
IndexTable rep -> IndexTable rep -> IndexTable rep
unionIndexTables IndexTable rep
indexTable
          ([IndexTable rep] -> IndexTable rep)
-> [IndexTable rep] -> IndexTable rep
forall a b. (a -> b) -> a -> b
$ ((VName, (ArrayName, [VName], [DimAccess rep]))
 -> Maybe (IndexTable rep))
-> [(VName, (ArrayName, [VName], [DimAccess rep]))]
-> [IndexTable rep]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
            ( (VName
 -> (ArrayName, [VName], [DimAccess rep]) -> Maybe (IndexTable rep))
-> (VName, (ArrayName, [VName], [DimAccess rep]))
-> Maybe (IndexTable rep)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((VName
  -> (ArrayName, [VName], [DimAccess rep]) -> Maybe (IndexTable rep))
 -> (VName, (ArrayName, [VName], [DimAccess rep]))
 -> Maybe (IndexTable rep))
-> (VName
    -> (ArrayName, [VName], [DimAccess rep]) -> Maybe (IndexTable rep))
-> (VName, (ArrayName, [VName], [DimAccess rep]))
-> Maybe (IndexTable rep)
forall a b. (a -> b) -> a -> b
$ \VName
_idx_expression (ArrayName
array_name, [VName]
patterns, [DimAccess rep]
dim_indices) ->
                IndexTable rep -> Maybe (IndexTable rep)
forall a. a -> Maybe a
Just (IndexTable rep -> Maybe (IndexTable rep))
-> ((Context rep, IndexTable rep) -> IndexTable rep)
-> (Context rep, IndexTable rep)
-> Maybe (IndexTable rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Context rep, IndexTable rep) -> IndexTable rep
forall a b. (a, b) -> b
snd ((Context rep, IndexTable rep) -> Maybe (IndexTable rep))
-> (Context rep, IndexTable rep) -> Maybe (IndexTable rep)
forall a b. (a -> b) -> a -> b
$
                  -- Should we use recContex instead of ctx''?
                  Context rep
-> [VName]
-> ArrayName
-> [DimAccess rep]
-> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> [VName]
-> ArrayName
-> [DimAccess rep]
-> (Context rep, IndexTable rep)
analyseIndex' Context rep
ctx'' [VName]
patterns ArrayName
array_name [DimAccess rep]
dim_indices
            )
          ([(VName, (ArrayName, [VName], [DimAccess rep]))]
 -> [IndexTable rep])
-> [(VName, (ArrayName, [VName], [DimAccess rep]))]
-> [IndexTable rep]
forall a b. (a -> b) -> a -> b
$ Map VName (ArrayName, [VName], [DimAccess rep])
-> [(VName, (ArrayName, [VName], [DimAccess rep]))]
forall k a. Map k a -> [(k, a)]
M.toList Map VName (ArrayName, [VName], [DimAccess rep])
slices_new

  -- 1. We do not want the returned context directly.
  --    however, we do want pat to map to the names what was hit in body.
  --    therefore we need to subtract the old context from the returned one,
  --    and discard all the keys within it.

  -- assignments :: M.Map VName (VariableInfo rep),
  let in_scope_dependencies_from_body :: Names
in_scope_dependencies_from_body =
        Context rep -> Map VName (VariableInfo rep) -> Names
forall {k} (rep :: k).
Context rep -> Map VName (VariableInfo rep) -> Names
rmOutOfScopeDeps Context rep
ctx'' (Map VName (VariableInfo rep) -> Names)
-> Map VName (VariableInfo rep) -> Names
forall a b. (a -> b) -> a -> b
$
          Map VName (VariableInfo rep)
-> Map VName (VariableInfo rep) -> Map VName (VariableInfo rep)
forall k a b. Ord k => Map k a -> Map k b -> Map k a
M.difference (Context rep -> Map VName (VariableInfo rep)
forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments Context rep
ctx'') (Context rep -> Map VName (VariableInfo rep)
forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments Context rep
recContext)

  -- 2. We are ONLY interested in the rhs of assignments (ie. the
  --    dependencies of pat :) )
  let ctx' :: Context rep
ctx' = (Context rep -> Context rep -> Context rep)
-> Context rep -> [Context rep] -> Context rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Context rep -> Context rep -> Context rep
forall k (rep :: k). Context rep -> Context rep -> Context rep
extend Context rep
ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
forall a b. (a -> b) -> a -> b
$ Names -> [Context rep]
concatVariableInfo Names
in_scope_dependencies_from_body -- . map snd $ M.toList varInfos
  -- 3. Now we have the correct context and result
  (Context rep
ctx' {parents = parents ctx, currentLevel = currentLevel ctx, slices = slices ctx}, IndexTable rep
slices_indices)
  where
    -- Extracts and merges `Names` in `VariableInfo`s, and makes a new VariableInfo. This
    -- MAY throw away needed information, but it was my best guess at a solution
    -- at the time of writing.
    concatVariableInfo :: Names -> [Context rep]
concatVariableInfo Names
dependencies =
      (VName -> Context rep) -> [VName] -> [Context rep]
forall a b. (a -> b) -> [a] -> [b]
map (\VName
pat -> VName -> VariableInfo rep -> Context rep
forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
oneContext VName
pat (Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx Names
dependencies)) [VName]
pats

    -- Context used for "recursion" into analyseStmsPrimitive
    recContext :: Context rep
recContext =
      Context rep
ctx
        { parents = parents ctx <> concatMap (\VName
pat -> [VName -> BodyType
body_constructor VName
pat]) pats,
          currentLevel = currentLevel ctx + 1
        }

    -- Recursively looks up dependencies, until they're in scope or empty set.
    rmOutOfScopeDeps :: Context rep -> M.Map VName (VariableInfo rep) -> Names
    rmOutOfScopeDeps :: forall {k} (rep :: k).
Context rep -> Map VName (VariableInfo rep) -> Names
rmOutOfScopeDeps Context rep
ctx' Map VName (VariableInfo rep)
new_assignments =
      let throwaway_assignments :: Map VName (VariableInfo rep)
throwaway_assignments = Context rep -> Map VName (VariableInfo rep)
forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments Context rep
ctx'
          local_assignments :: Map VName (VariableInfo rep)
local_assignments = Context rep -> Map VName (VariableInfo rep)
forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments Context rep
ctx
          f :: Names -> VName -> VariableInfo rep -> Names
f Names
result VName
a VariableInfo rep
var_info =
            -- if the VName of the assignment exists in the context, we are good
            if VName
a VName -> Map VName (VariableInfo rep) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName (VariableInfo rep)
local_assignments
              then Names
result Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> VName -> Names
oneName VName
a
              else -- Otherwise, recurse on its dependencies;
              -- 0. Add dependencies in ctx to result

                let ([VName]
deps_in_ctx, [VName]
deps_not_in_ctx) =
                      (VName -> Bool) -> [VName] -> ([VName], [VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (VName -> Map VName (VariableInfo rep) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName (VariableInfo rep)
local_assignments) ([VName] -> ([VName], [VName])) -> [VName] -> ([VName], [VName])
forall a b. (a -> b) -> a -> b
$
                        Names -> [VName]
namesToList (VariableInfo rep -> Names
forall {k} (rep :: k). VariableInfo rep -> Names
deps VariableInfo rep
var_info)
                    deps_not_in_ctx' :: Map VName (VariableInfo rep)
deps_not_in_ctx' =
                      [(VName, VariableInfo rep)] -> Map VName (VariableInfo rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VariableInfo rep)] -> Map VName (VariableInfo rep))
-> [(VName, VariableInfo rep)] -> Map VName (VariableInfo rep)
forall a b. (a -> b) -> a -> b
$
                        (VName -> Maybe (VName, VariableInfo rep))
-> [VName] -> [(VName, VariableInfo rep)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
                          (\VName
d -> (VName
d,) (VariableInfo rep -> (VName, VariableInfo rep))
-> Maybe (VariableInfo rep) -> Maybe (VName, VariableInfo rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName (VariableInfo rep) -> Maybe (VariableInfo rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
d Map VName (VariableInfo rep)
throwaway_assignments)
                          [VName]
deps_not_in_ctx
                 in Names
result
                      Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList [VName]
deps_in_ctx
                      Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Context rep -> Map VName (VariableInfo rep) -> Names
forall {k} (rep :: k).
Context rep -> Map VName (VariableInfo rep) -> Names
rmOutOfScopeDeps Context rep
ctx' Map VName (VariableInfo rep)
deps_not_in_ctx'
       in (Names -> VName -> VariableInfo rep -> Names)
-> Names -> Map VName (VariableInfo rep) -> Names
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey Names -> VName -> VariableInfo rep -> Names
forall {k} {rep :: k}. Names -> VName -> VariableInfo rep -> Names
f Names
forall a. Monoid a => a
mempty Map VName (VariableInfo rep)
new_assignments

-- | Analyse a rep statement and return the updated context and array index
-- descriptors.
analyseStm :: (Analyse rep) => Context rep -> Stm rep -> (Context rep, IndexTable rep)
analyseStm :: forall rep.
Analyse rep =>
Context rep -> Stm rep -> (Context rep, IndexTable rep)
analyseStm Context rep
ctx (Let Pat (LetDec rep)
pats StmAux (ExpDec rep)
_ Exp rep
e) = do
  -- Get the name of the first element in a pattern
  let pattern_names :: [VName]
pattern_names = (PatElem (LetDec rep) -> VName)
-> [PatElem (LetDec rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName ([PatElem (LetDec rep)] -> [VName])
-> [PatElem (LetDec rep)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec rep)
pats

  -- Construct the result and Context from the subexpression. If the
  -- subexpression is a body, we recurse into it.
  case Exp rep
e of
    BasicOp (Index VName
name (Slice [DimIndex SubExp]
dim_subexp)) ->
      Context rep
-> [VName]
-> VName
-> [DimIndex SubExp]
-> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> [VName]
-> VName
-> [DimIndex SubExp]
-> (Context rep, IndexTable rep)
analyseIndex Context rep
ctx [VName]
pattern_names VName
name [DimIndex SubExp]
dim_subexp
    BasicOp (Update Safety
_ VName
name (Slice [DimIndex SubExp]
dim_subexp) SubExp
_subexp) ->
      Context rep
-> [VName]
-> VName
-> [DimIndex SubExp]
-> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> [VName]
-> VName
-> [DimIndex SubExp]
-> (Context rep, IndexTable rep)
analyseIndex Context rep
ctx [VName]
pattern_names VName
name [DimIndex SubExp]
dim_subexp
    BasicOp BasicOp
op ->
      Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
analyseBasicOp Context rep
ctx BasicOp
op [VName]
pattern_names
    Match [SubExp]
conds [Case (Body rep)]
cases Body rep
default_body MatchDec (BranchType rep)
_ ->
      Context rep
-> [VName]
-> Body rep
-> [Body rep]
-> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep
-> [VName]
-> Body rep
-> [Body rep]
-> (Context rep, IndexTable rep)
analyseMatch Context rep
ctx' [VName]
pattern_names Body rep
default_body ([Body rep] -> (Context rep, IndexTable rep))
-> [Body rep] -> (Context rep, IndexTable rep)
forall a b. (a -> b) -> a -> b
$ (Case (Body rep) -> Body rep) -> [Case (Body rep)] -> [Body rep]
forall a b. (a -> b) -> [a] -> [b]
map Case (Body rep) -> Body rep
forall body. Case body -> body
caseBody [Case (Body rep)]
cases
      where
        ctx' :: Context rep
ctx' =
          Context rep -> VariableInfo rep -> [VName] -> Context rep
forall {k} (rep :: k).
Context rep -> VariableInfo rep -> [VName] -> Context rep
contextFromNames Context rep
ctx (Context rep -> VariableInfo rep
forall {k} (rep :: k). Context rep -> VariableInfo rep
varInfoZeroDeps Context rep
ctx) ([VName] -> Context rep) -> [VName] -> Context rep
forall a b. (a -> b) -> a -> b
$
            (SubExp -> [VName]) -> [SubExp] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Names -> [VName]
namesToList (Names -> [VName]) -> (SubExp -> Names) -> SubExp -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn) [SubExp]
conds
    Loop [(FParam rep, SubExp)]
bindings LoopForm
loop Body rep
body ->
      Context rep
-> [(FParam rep, SubExp)]
-> LoopForm
-> Body rep
-> [VName]
-> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep
-> [(FParam rep, SubExp)]
-> LoopForm
-> Body rep
-> [VName]
-> (Context rep, IndexTable rep)
analyseLoop Context rep
ctx [(FParam rep, SubExp)]
bindings LoopForm
loop Body rep
body [VName]
pattern_names
    Apply Name
_name [(SubExp, Diet)]
diets [(RetType rep, RetAls)]
_ (Safety, SrcLoc, [SrcLoc])
_ ->
      Context rep
-> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
analyseApply Context rep
ctx [VName]
pattern_names [(SubExp, Diet)]
diets
    WithAcc [WithAccInput rep]
_ Lambda rep
_ ->
      (Context rep
ctx, IndexTable rep
forall a. Monoid a => a
mempty) -- ignored
    Op Op rep
op ->
      Op rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Op rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseOp Op rep
op Context rep
ctx [VName]
pattern_names

-- If left, this is just a regular index. If right, a slice happened.
getIndexDependencies :: Context rep -> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
getIndexDependencies :: forall {k} (rep :: k).
Context rep
-> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
getIndexDependencies Context rep
ctx [DimIndex SubExp]
dims =
  (Either [DimAccess rep] [DimAccess rep], Int)
-> Either [DimAccess rep] [DimAccess rep]
forall a b. (a, b) -> a
fst ((Either [DimAccess rep] [DimAccess rep], Int)
 -> Either [DimAccess rep] [DimAccess rep])
-> (Either [DimAccess rep] [DimAccess rep], Int)
-> Either [DimAccess rep] [DimAccess rep]
forall a b. (a -> b) -> a -> b
$
    (DimIndex SubExp
 -> (Either [DimAccess rep] [DimAccess rep], Int)
 -> (Either [DimAccess rep] [DimAccess rep], Int))
-> (Either [DimAccess rep] [DimAccess rep], Int)
-> [DimIndex SubExp]
-> (Either [DimAccess rep] [DimAccess rep], Int)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
      ( \DimIndex SubExp
idx (Either [DimAccess rep] [DimAccess rep]
a, Int
i) ->
          ( ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
-> ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
-> Either [DimAccess rep] [DimAccess rep]
-> Either [DimAccess rep] [DimAccess rep]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (DimIndex SubExp
-> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
matchDimIndex DimIndex SubExp
idx) (([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
-> ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
-> Either [DimAccess rep] [DimAccess rep]
-> Either [DimAccess rep] [DimAccess rep]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
forall a b. b -> Either a b
Right [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
forall a b. b -> Either a b
Right (Either [DimAccess rep] [DimAccess rep]
 -> Either [DimAccess rep] [DimAccess rep])
-> ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
-> [DimAccess rep]
-> Either [DimAccess rep] [DimAccess rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DimIndex SubExp
-> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
matchDimIndex DimIndex SubExp
idx) Either [DimAccess rep] [DimAccess rep]
a,
            Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
          )
      )
      ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
forall a b. a -> Either a b
Left [], [DimIndex SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex SubExp]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
      [DimIndex SubExp]
dims
  where
    matchDimIndex :: DimIndex SubExp
-> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
matchDimIndex (DimFix SubExp
subExpression) [DimAccess rep]
accumulator =
      [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
forall a b. a -> Either a b
Left ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
-> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
forall a b. (a -> b) -> a -> b
$ Context rep -> SubExp -> DimAccess rep
forall {k} (rep :: k). Context rep -> SubExp -> DimAccess rep
consolidate Context rep
ctx SubExp
subExpression DimAccess rep -> [DimAccess rep] -> [DimAccess rep]
forall a. a -> [a] -> [a]
: [DimAccess rep]
accumulator
    -- If we encounter a DimSlice, add it to a map of `DimSlice`s and check
    -- result later.
    matchDimIndex (DimSlice SubExp
offset SubExp
num_elems SubExp
stride) [DimAccess rep]
accumulator =
      -- We assume that a slice is iterated sequentially, so we have to
      -- create a fake dependency for the slice.
      let dimAccess' :: DimAccess rep
dimAccess' = Map VName Dependency -> Maybe VName -> DimAccess rep
forall {k} (rep :: k).
Map VName Dependency -> Maybe VName -> DimAccess rep
DimAccess (VName -> Dependency -> Map VName Dependency
forall k a. k -> a -> Map k a
M.singleton (Name -> Int -> VName
VName Name
"slice" Int
0) (Dependency -> Map VName Dependency)
-> Dependency -> Map VName Dependency
forall a b. (a -> b) -> a -> b
$ Int -> VarType -> Dependency
Dependency (Context rep -> Int
forall {k} (rep :: k). Context rep -> Int
currentLevel Context rep
ctx) VarType
LoopVar) (VName -> Maybe VName
forall a. a -> Maybe a
Just (VName -> Maybe VName) -> VName -> Maybe VName
forall a b. (a -> b) -> a -> b
$ Name -> Int -> VName
VName Name
"slice" Int
0)
          cons :: SubExp -> DimAccess rep
cons = Context rep -> SubExp -> DimAccess rep
forall {k} (rep :: k). Context rep -> SubExp -> DimAccess rep
consolidate Context rep
ctx
          dimAccess :: DimAccess rep
dimAccess = DimAccess rep
forall k (rep :: k). DimAccess rep
dimAccess' DimAccess rep -> DimAccess rep -> DimAccess rep
forall a. Semigroup a => a -> a -> a
<> SubExp -> DimAccess rep
cons SubExp
offset DimAccess rep -> DimAccess rep -> DimAccess rep
forall a. Semigroup a => a -> a -> a
<> SubExp -> DimAccess rep
cons SubExp
num_elems DimAccess rep -> DimAccess rep -> DimAccess rep
forall a. Semigroup a => a -> a -> a
<> SubExp -> DimAccess rep
cons SubExp
stride
       in [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
forall a b. b -> Either a b
Right ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
-> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
forall a b. (a -> b) -> a -> b
$ DimAccess rep
dimAccess DimAccess rep -> [DimAccess rep] -> [DimAccess rep]
forall a. a -> [a] -> [a]
: [DimAccess rep]
accumulator

-- | Gets the dependencies of each dimension and either returns a result, or
-- adds a slice to the context.
analyseIndex :: Context rep -> [VName] -> VName -> [DimIndex SubExp] -> (Context rep, IndexTable rep)
analyseIndex :: forall {k} (rep :: k).
Context rep
-> [VName]
-> VName
-> [DimIndex SubExp]
-> (Context rep, IndexTable rep)
analyseIndex Context rep
ctx [VName]
pats VName
arr_name [DimIndex SubExp]
dim_indices =
  -- Get the dependendencies of each dimension
  let dependencies :: Either [DimAccess rep] [DimAccess rep]
dependencies = Context rep
-> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
forall {k} (rep :: k).
Context rep
-> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
getIndexDependencies Context rep
ctx [DimIndex SubExp]
dim_indices
      -- Extend the current context with current pattern(s) and its deps
      ctx' :: Context rep
ctx' = Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
forall {k} (rep :: k).
Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
analyseIndexContextFromIndices Context rep
ctx [DimIndex SubExp]
dim_indices [VName]
pats

      -- The bodytype(s) are used in the result construction
      array_name' :: ArrayName
array_name' =
        -- For now, we assume the array is in row-major-order, hence the
        -- identity permutation. In the future, we might want to infer its
        -- layout, for example, if the array is the result of a transposition.
        let layout :: [Int]
layout = [Int
0 .. [DimIndex SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex SubExp]
dim_indices Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
         in -- 2. If the arrayname was not in assignments, it was not an immediately
            --    allocated array.
            ArrayName -> Maybe ArrayName -> ArrayName
forall a. a -> Maybe a -> a
fromMaybe (VName
arr_name, [], [Int]
layout)
              -- 1. Maybe find the array name, and the "stack" of body types that the
              -- array was allocated in.
              (Maybe ArrayName -> ArrayName)
-> ([ArrayName] -> Maybe ArrayName) -> [ArrayName] -> ArrayName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ArrayName -> Bool) -> [ArrayName] -> Maybe ArrayName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
L.find (\(VName
n, [BodyType]
_, [Int]
_) -> VName
n VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
arr_name)
              -- 0. Get the "stack" of bodytypes for each assignment
              ([ArrayName] -> ArrayName) -> [ArrayName] -> ArrayName
forall a b. (a -> b) -> a -> b
$ ((VName, VariableInfo rep) -> ArrayName)
-> [(VName, VariableInfo rep)] -> [ArrayName]
forall a b. (a -> b) -> [a] -> [b]
map (\(VName
n, VariableInfo rep
vi) -> (VName
n, VariableInfo rep -> [BodyType]
forall {k} (rep :: k). VariableInfo rep -> [BodyType]
parents_nest VariableInfo rep
vi, [Int]
layout)) (Map VName (VariableInfo rep) -> [(VName, VariableInfo rep)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName (VariableInfo rep) -> [(VName, VariableInfo rep)])
-> Map VName (VariableInfo rep) -> [(VName, VariableInfo rep)]
forall a b. (a -> b) -> a -> b
$ Context rep -> Map VName (VariableInfo rep)
forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments Context rep
ctx')
   in ([DimAccess rep] -> (Context rep, IndexTable rep))
-> ([DimAccess rep] -> (Context rep, IndexTable rep))
-> Either [DimAccess rep] [DimAccess rep]
-> (Context rep, IndexTable rep)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Context rep
-> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
index Context rep
ctx' ArrayName
array_name') (Context rep
-> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
slice Context rep
ctx' ArrayName
array_name') Either [DimAccess rep] [DimAccess rep]
dependencies
  where
    slice :: Context rep -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    slice :: forall {k} (rep :: k).
Context rep
-> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
slice Context rep
context ArrayName
array_name [DimAccess rep]
dims =
      (Context rep
context {slices = M.insert (head pats) (array_name, pats, dims) $ slices context}, IndexTable rep
forall a. Monoid a => a
mempty)

    index :: Context rep -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    index :: forall {k} (rep :: k).
Context rep
-> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
index Context rep
context array_name :: ArrayName
array_name@(VName
name, [BodyType]
_, [Int]
_) [DimAccess rep]
dim_access =
      -- If the arrayname is a `DimSlice` we want to fixup the access
      case VName
-> Map VName (ArrayName, [VName], [DimAccess rep])
-> Maybe (ArrayName, [VName], [DimAccess rep])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName (ArrayName, [VName], [DimAccess rep])
 -> Maybe (ArrayName, [VName], [DimAccess rep]))
-> Map VName (ArrayName, [VName], [DimAccess rep])
-> Maybe (ArrayName, [VName], [DimAccess rep])
forall a b. (a -> b) -> a -> b
$ Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
forall {k} (rep :: k).
Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
slices Context rep
context of
        Maybe (ArrayName, [VName], [DimAccess rep])
Nothing -> Context rep
-> [VName]
-> ArrayName
-> [DimAccess rep]
-> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> [VName]
-> ArrayName
-> [DimAccess rep]
-> (Context rep, IndexTable rep)
analyseIndex' Context rep
context [VName]
pats ArrayName
array_name [DimAccess rep]
dim_access
        Just (ArrayName
arr_name', [VName]
pats', [DimAccess rep]
slice_access) ->
          Context rep
-> [VName]
-> ArrayName
-> [DimAccess rep]
-> (Context rep, IndexTable rep)
forall {k} (rep :: k).
Context rep
-> [VName]
-> ArrayName
-> [DimAccess rep]
-> (Context rep, IndexTable rep)
analyseIndex'
            Context rep
context
            [VName]
pats'
            ArrayName
arr_name'
            ([DimAccess rep] -> [DimAccess rep]
forall a. HasCallStack => [a] -> [a]
init [DimAccess rep]
slice_access [DimAccess rep] -> [DimAccess rep] -> [DimAccess rep]
forall a. [a] -> [a] -> [a]
++ [[DimAccess rep] -> DimAccess rep
forall a. HasCallStack => [a] -> a
head [DimAccess rep]
dim_access DimAccess rep -> DimAccess rep -> DimAccess rep
forall a. Semigroup a => a -> a -> a
<> [DimAccess rep] -> DimAccess rep
forall a. HasCallStack => [a] -> a
last [DimAccess rep]
slice_access] [DimAccess rep] -> [DimAccess rep] -> [DimAccess rep]
forall a. [a] -> [a] -> [a]
++ Int -> [DimAccess rep] -> [DimAccess rep]
forall a. Int -> [a] -> [a]
drop Int
1 [DimAccess rep]
dim_access)

analyseIndexContextFromIndices :: Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
analyseIndexContextFromIndices :: forall {k} (rep :: k).
Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
analyseIndexContextFromIndices Context rep
ctx [DimIndex SubExp]
dim_accesses [VName]
pats =
  let subexprs :: [VName]
subexprs =
        (DimIndex SubExp -> Maybe VName) -> [DimIndex SubExp] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
          ( \case
              DimFix (Var VName
v) -> VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
              DimFix (Constant PrimValue
_) -> Maybe VName
forall a. Maybe a
Nothing
              DimSlice SubExp
_offs SubExp
_n SubExp
_stride -> Maybe VName
forall a. Maybe a
Nothing
          )
          [DimIndex SubExp]
dim_accesses

      -- Add each non-constant DimIndex as a dependency to the index expression
      var_info :: VariableInfo rep
var_info = Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ [VName] -> Names
namesFromList [VName]
subexprs
   in -- Extend context with the dependencies index expression
      (Context rep -> Context rep -> Context rep)
-> Context rep -> [Context rep] -> Context rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Context rep -> Context rep -> Context rep
forall k (rep :: k). Context rep -> Context rep -> Context rep
extend Context rep
ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
forall a b. (a -> b) -> a -> b
$ (VName -> Context rep) -> [VName] -> [Context rep]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> VariableInfo rep -> Context rep
forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
`oneContext` VariableInfo rep
var_info) [VName]
pats

analyseIndex' ::
  Context rep ->
  [VName] ->
  ArrayName ->
  [DimAccess rep] ->
  (Context rep, IndexTable rep)
analyseIndex' :: forall {k} (rep :: k).
Context rep
-> [VName]
-> ArrayName
-> [DimAccess rep]
-> (Context rep, IndexTable rep)
analyseIndex' Context rep
ctx [VName]
_ ArrayName
_ [] = (Context rep
ctx, IndexTable rep
forall a. Monoid a => a
mempty)
analyseIndex' Context rep
ctx [VName]
_ ArrayName
_ [DimAccess rep
_] = (Context rep
ctx, IndexTable rep
forall a. Monoid a => a
mempty)
analyseIndex' Context rep
ctx [VName]
pats ArrayName
arr_name [DimAccess rep]
dim_accesses =
  -- Get the name of all segmaps in the current "callstack"
  let segmaps :: [SegOpName]
segmaps = Context rep -> [SegOpName]
forall {k} (rep :: k). Context rep -> [SegOpName]
allSegMap Context rep
ctx
      idx_expr_name :: [VName]
idx_expr_name = [VName]
pats --                                                IndexExprName
      -- For each pattern, create a mapping to the dimensional indices
      map_ixd_expr :: [Map VName [DimAccess rep]]
map_ixd_expr = (VName -> Map VName [DimAccess rep])
-> [VName] -> [Map VName [DimAccess rep]]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> [DimAccess rep] -> Map VName [DimAccess rep]
forall k a. k -> a -> Map k a
`M.singleton` [DimAccess rep]
dim_accesses) [VName]
idx_expr_name --       IndexExprName |-> [DimAccess]
      -- For each pattern -> [DimAccess] mapping, create a mapping from the array
      -- name that was indexed.
      map_array :: [Map ArrayName (Map VName [DimAccess rep])]
map_array = (Map VName [DimAccess rep]
 -> Map ArrayName (Map VName [DimAccess rep]))
-> [Map VName [DimAccess rep]]
-> [Map ArrayName (Map VName [DimAccess rep])]
forall a b. (a -> b) -> [a] -> [b]
map (ArrayName
-> Map VName [DimAccess rep]
-> Map ArrayName (Map VName [DimAccess rep])
forall k a. k -> a -> Map k a
M.singleton ArrayName
arr_name) [Map VName [DimAccess rep]]
map_ixd_expr --   ArrayName |-> IndexExprName |-> [DimAccess]
      -- ∀ (arr_name -> IdxExp -> [DimAccess]) mappings, create a mapping from all
      -- segmaps in current callstack (segThread & segGroups alike).
      results :: [IndexTable rep]
results = (Map ArrayName (Map VName [DimAccess rep]) -> [IndexTable rep])
-> [Map ArrayName (Map VName [DimAccess rep])] -> [IndexTable rep]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Map ArrayName (Map VName [DimAccess rep])
ma -> (SegOpName -> IndexTable rep) -> [SegOpName] -> [IndexTable rep]
forall a b. (a -> b) -> [a] -> [b]
map (SegOpName
-> Map ArrayName (Map VName [DimAccess rep]) -> IndexTable rep
forall k a. k -> a -> Map k a
`M.singleton` Map ArrayName (Map VName [DimAccess rep])
ma) [SegOpName]
segmaps) [Map ArrayName (Map VName [DimAccess rep])]
map_array

      res :: IndexTable rep
res = (IndexTable rep -> IndexTable rep -> IndexTable rep)
-> IndexTable rep -> [IndexTable rep] -> IndexTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IndexTable rep -> IndexTable rep -> IndexTable rep
forall {k} (rep :: k).
IndexTable rep -> IndexTable rep -> IndexTable rep
unionIndexTables IndexTable rep
forall a. Monoid a => a
mempty [IndexTable rep]
results
   in (Context rep
ctx, IndexTable rep
res)

analyseBasicOp :: Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
analyseBasicOp :: forall {k} (rep :: k).
Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
analyseBasicOp Context rep
ctx BasicOp
expression [VName]
pats =
  -- Construct a VariableInfo from the subexpressions
  let ctx_val :: VariableInfo rep
ctx_val = case BasicOp
expression of
        SubExp SubExp
se -> SubExp -> VariableInfo rep
forall {k} {rep :: k}. SubExp -> VariableInfo rep
varInfoFromSubExp SubExp
se
        Opaque OpaqueOp
_ SubExp
se -> SubExp -> VariableInfo rep
forall {k} {rep :: k}. SubExp -> VariableInfo rep
varInfoFromSubExp SubExp
se
        ArrayVal [PrimValue]
_ PrimType
_ -> (Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx Names
forall a. Monoid a => a
mempty) {variableType = ConstType}
        ArrayLit [SubExp]
ses Type
_t -> Names -> [SubExp] -> VariableInfo rep
concatVariableInfos Names
forall a. Monoid a => a
mempty [SubExp]
ses
        UnOp UnOp
_ SubExp
se -> SubExp -> VariableInfo rep
forall {k} {rep :: k}. SubExp -> VariableInfo rep
varInfoFromSubExp SubExp
se
        BinOp BinOp
_ SubExp
lsubexp SubExp
rsubexp -> Names -> [SubExp] -> VariableInfo rep
concatVariableInfos Names
forall a. Monoid a => a
mempty [SubExp
lsubexp, SubExp
rsubexp]
        CmpOp CmpOp
_ SubExp
lsubexp SubExp
rsubexp -> Names -> [SubExp] -> VariableInfo rep
concatVariableInfos Names
forall a. Monoid a => a
mempty [SubExp
lsubexp, SubExp
rsubexp]
        ConvOp ConvOp
_ SubExp
se -> SubExp -> VariableInfo rep
forall {k} {rep :: k}. SubExp -> VariableInfo rep
varInfoFromSubExp SubExp
se
        Assert SubExp
se ErrorMsg SubExp
_ (SrcLoc, [SrcLoc])
_ -> SubExp -> VariableInfo rep
forall {k} {rep :: k}. SubExp -> VariableInfo rep
varInfoFromSubExp SubExp
se
        Index VName
name Slice SubExp
_ ->
          [Char] -> VariableInfo rep
forall a. HasCallStack => [Char] -> a
error ([Char] -> VariableInfo rep) -> [Char] -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ [Char]
"unhandled: Index (This should NEVER happen) into " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
name
        Update Safety
_ VName
name Slice SubExp
_slice SubExp
_subexp ->
          [Char] -> VariableInfo rep
forall a. HasCallStack => [Char] -> a
error ([Char] -> VariableInfo rep) -> [Char] -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ [Char]
"unhandled: Update (This should NEVER happen) onto " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
name
        -- Technically, do we need this case?
        Concat Int
_ NonEmpty VName
_ SubExp
length_subexp -> SubExp -> VariableInfo rep
forall {k} {rep :: k}. SubExp -> VariableInfo rep
varInfoFromSubExp SubExp
length_subexp
        Manifest [Int]
_dim VName
name -> Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
name
        Iota SubExp
end SubExp
start SubExp
stride IntType
_ -> Names -> [SubExp] -> VariableInfo rep
concatVariableInfos Names
forall a. Monoid a => a
mempty [SubExp
end, SubExp
start, SubExp
stride]
        Replicate (Shape [SubExp]
shape) SubExp
value' -> Names -> [SubExp] -> VariableInfo rep
concatVariableInfos Names
forall a. Monoid a => a
mempty (SubExp
value' SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
shape)
        Scratch PrimType
_ [SubExp]
sers -> Names -> [SubExp] -> VariableInfo rep
concatVariableInfos Names
forall a. Monoid a => a
mempty [SubExp]
sers
        Reshape ReshapeKind
_ (Shape [SubExp]
shape_subexp) VName
name -> Names -> [SubExp] -> VariableInfo rep
concatVariableInfos (VName -> Names
oneName VName
name) [SubExp]
shape_subexp
        Rearrange [Int]
_ VName
name -> Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
name
        UpdateAcc Safety
_ VName
name [SubExp]
lsubexprs [SubExp]
rsubexprs ->
          Names -> [SubExp] -> VariableInfo rep
concatVariableInfos (VName -> Names
oneName VName
name) ([SubExp]
lsubexprs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
rsubexprs)
        FlatIndex VName
name FlatSlice SubExp
_ -> Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
name
        FlatUpdate VName
name FlatSlice SubExp
_ VName
source -> Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ [VName] -> Names
namesFromList [VName
name, VName
source]
      ctx' :: Context rep
ctx' = (Context rep -> Context rep -> Context rep)
-> Context rep -> [Context rep] -> Context rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Context rep -> Context rep -> Context rep
forall k (rep :: k). Context rep -> Context rep -> Context rep
extend Context rep
ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
forall a b. (a -> b) -> a -> b
$ (VName -> Context rep) -> [VName] -> [Context rep]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> VariableInfo rep -> Context rep
forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
`oneContext` VariableInfo rep
ctx_val) [VName]
pats
   in (Context rep
ctx', IndexTable rep
forall a. Monoid a => a
mempty)
  where
    concatVariableInfos :: Names -> [SubExp] -> VariableInfo rep
concatVariableInfos Names
ne [SubExp]
nn =
      Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names
ne Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ((SubExp -> Names) -> [SubExp] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Context rep -> SubExp -> Names
forall {k} (rep :: k). [VName] -> Context rep -> SubExp -> Names
analyseSubExp [VName]
pats Context rep
ctx) [SubExp]
nn))

    varInfoFromSubExp :: SubExp -> VariableInfo rep
varInfoFromSubExp (Constant PrimValue
_) = (Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx Names
forall a. Monoid a => a
mempty) {variableType = ConstType}
    varInfoFromSubExp (Var VName
v) =
      case VName -> Map VName (VariableInfo rep) -> Maybe (VariableInfo rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Context rep -> Map VName (VariableInfo rep)
forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments Context rep
ctx) of
        Just VariableInfo rep
_ -> (Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
v) {variableType = Variable}
        Maybe (VariableInfo rep)
Nothing -> (Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx Names
forall a. Monoid a => a
mempty) {variableType = Variable} -- Means a global.

analyseMatch :: (Analyse rep) => Context rep -> [VName] -> Body rep -> [Body rep] -> (Context rep, IndexTable rep)
analyseMatch :: forall rep.
Analyse rep =>
Context rep
-> [VName]
-> Body rep
-> [Body rep]
-> (Context rep, IndexTable rep)
analyseMatch Context rep
ctx [VName]
pats Body rep
body [Body rep]
parents =
  let ctx'' :: Context rep
ctx'' = Context rep
ctx {currentLevel = currentLevel ctx - 1}
   in ((Context rep, IndexTable rep)
 -> Body rep -> (Context rep, IndexTable rep))
-> (Context rep, IndexTable rep)
-> [Body rep]
-> (Context rep, IndexTable rep)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
        ( \(Context rep
ctx', IndexTable rep
res) Body rep
b ->
            -- This Little Maneuver's Gonna Cost Us 51 Years
            (Context rep -> Context rep)
-> (IndexTable rep -> IndexTable rep)
-> (Context rep, IndexTable rep)
-> (Context rep, IndexTable rep)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Context rep -> Context rep
forall {k} {rep :: k}. Context rep -> Context rep
constLevel (IndexTable rep -> IndexTable rep -> IndexTable rep
forall {k} (rep :: k).
IndexTable rep -> IndexTable rep -> IndexTable rep
unionIndexTables IndexTable rep
res)
              ((Context rep, IndexTable rep) -> (Context rep, IndexTable rep))
-> (Stms rep -> (Context rep, IndexTable rep))
-> Stms rep
-> (Context rep, IndexTable rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Context rep
-> (VName -> BodyType)
-> [VName]
-> [Stm rep]
-> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep
-> (VName -> BodyType)
-> [VName]
-> [Stm rep]
-> (Context rep, IndexTable rep)
analyseStms Context rep
ctx' VName -> BodyType
CondBodyName [VName]
pats
              ([Stm rep] -> (Context rep, IndexTable rep))
-> (Stms rep -> [Stm rep])
-> Stms rep
-> (Context rep, IndexTable rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList
              (Stms rep -> (Context rep, IndexTable rep))
-> Stms rep -> (Context rep, IndexTable rep)
forall a b. (a -> b) -> a -> b
$ Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
b
        )
        (Context rep
ctx'', IndexTable rep
forall a. Monoid a => a
mempty)
        (Body rep
body Body rep -> [Body rep] -> [Body rep]
forall a. a -> [a] -> [a]
: [Body rep]
parents)
  where
    constLevel :: Context rep -> Context rep
constLevel Context rep
context = Context rep
context {currentLevel = currentLevel ctx - 1}

analyseLoop :: (Analyse rep) => Context rep -> [(FParam rep, SubExp)] -> LoopForm -> Body rep -> [VName] -> (Context rep, IndexTable rep)
analyseLoop :: forall rep.
Analyse rep =>
Context rep
-> [(FParam rep, SubExp)]
-> LoopForm
-> Body rep
-> [VName]
-> (Context rep, IndexTable rep)
analyseLoop Context rep
ctx [(FParam rep, SubExp)]
bindings LoopForm
loop Body rep
body [VName]
pats = do
  let next_level :: Int
next_level = Context rep -> Int
forall {k} (rep :: k). Context rep -> Int
currentLevel Context rep
ctx
  let ctx'' :: Context rep
ctx'' = Context rep
ctx {currentLevel = next_level}
  let ctx' :: Context rep
ctx' =
        Context rep -> VariableInfo rep -> [VName] -> Context rep
forall {k} (rep :: k).
Context rep -> VariableInfo rep -> [VName] -> Context rep
contextFromNames Context rep
ctx'' ((Context rep -> VariableInfo rep
forall {k} (rep :: k). Context rep -> VariableInfo rep
varInfoZeroDeps Context rep
ctx) {variableType = LoopVar}) ([VName] -> Context rep) -> [VName] -> Context rep
forall a b. (a -> b) -> a -> b
$
          case LoopForm
loop of
            WhileLoop VName
iv -> VName
iv VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: ((FParam rep, SubExp) -> VName)
-> [(FParam rep, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (FParam rep -> VName
forall dec. Param dec -> VName
paramName (FParam rep -> VName)
-> ((FParam rep, SubExp) -> FParam rep)
-> (FParam rep, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> FParam rep
forall a b. (a, b) -> a
fst) [(FParam rep, SubExp)]
bindings
            ForLoop VName
iv IntType
_ SubExp
_ -> VName
iv VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: ((FParam rep, SubExp) -> VName)
-> [(FParam rep, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (FParam rep -> VName
forall dec. Param dec -> VName
paramName (FParam rep -> VName)
-> ((FParam rep, SubExp) -> FParam rep)
-> (FParam rep, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> FParam rep
forall a b. (a, b) -> a
fst) [(FParam rep, SubExp)]
bindings

  -- Extend context with the loop expression
  Context rep
-> (VName -> BodyType)
-> [VName]
-> [Stm rep]
-> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep
-> (VName -> BodyType)
-> [VName]
-> [Stm rep]
-> (Context rep, IndexTable rep)
analyseStms Context rep
ctx' VName -> BodyType
LoopBodyName [VName]
pats ([Stm rep] -> (Context rep, IndexTable rep))
-> [Stm rep] -> (Context rep, IndexTable rep)
forall a b. (a -> b) -> a -> b
$ Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms rep -> [Stm rep]) -> Stms rep -> [Stm rep]
forall a b. (a -> b) -> a -> b
$ Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
body

analyseApply :: Context rep -> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
analyseApply :: forall {k} (rep :: k).
Context rep
-> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
analyseApply Context rep
ctx [VName]
pats [(SubExp, Diet)]
diets =
  ( (Context rep -> Context rep -> Context rep)
-> Context rep -> [Context rep] -> Context rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Context rep -> Context rep -> Context rep
forall k (rep :: k). Context rep -> Context rep -> Context rep
extend Context rep
ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
forall a b. (a -> b) -> a -> b
$ (VName -> Context rep) -> [VName] -> [Context rep]
forall a b. (a -> b) -> [a] -> [b]
map (\VName
pat -> VName -> VariableInfo rep -> Context rep
forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
oneContext VName
pat (VariableInfo rep -> Context rep)
-> VariableInfo rep -> Context rep
forall a b. (a -> b) -> a -> b
$ Context rep -> Names -> VariableInfo rep
forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
varInfoFromNames Context rep
ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
forall a b. (a -> b) -> a -> b
$ [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> [Names] -> Names
forall a b. (a -> b) -> a -> b
$ ((SubExp, Diet) -> Names) -> [(SubExp, Diet)] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn (SubExp -> Names)
-> ((SubExp, Diet) -> SubExp) -> (SubExp, Diet) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst) [(SubExp, Diet)]
diets) [VName]
pats,
    IndexTable rep
forall a. Monoid a => a
mempty
  )

segOpType :: SegOp lvl rep -> VName -> SegOpName
segOpType :: forall lvl rep. SegOp lvl rep -> VName -> SegOpName
segOpType (SegMap {}) = VName -> SegOpName
SegmentedMap
segOpType (SegRed {}) = VName -> SegOpName
SegmentedRed
segOpType (SegScan {}) = VName -> SegOpName
SegmentedScan
segOpType (SegHist {}) = VName -> SegOpName
SegmentedHist

analyseSegOp :: (Analyse rep) => SegOp lvl rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSegOp :: forall rep lvl.
Analyse rep =>
SegOp lvl rep
-> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSegOp SegOp lvl rep
op Context rep
ctx [VName]
pats =
  let next_level :: Int
next_level = Context rep -> Int
forall {k} (rep :: k). Context rep -> Int
currentLevel Context rep
ctx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [(VName, SubExp)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SegSpace -> [(VName, SubExp)]
unSegSpace (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      ctx' :: Context rep
ctx' = Context rep
ctx {currentLevel = next_level}
      segspace_context :: Context rep
segspace_context =
        (Context rep -> Context rep -> Context rep)
-> Context rep -> [Context rep] -> Context rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Context rep -> Context rep -> Context rep
forall k (rep :: k). Context rep -> Context rep -> Context rep
extend Context rep
ctx'
          ([Context rep] -> Context rep)
-> (SegSpace -> [Context rep]) -> SegSpace -> Context rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, Int) -> Context rep) -> [(VName, Int)] -> [Context rep]
forall a b. (a -> b) -> [a] -> [b]
map (\(VName
n, Int
i) -> VName -> VariableInfo rep -> Context rep
forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
oneContext VName
n (VariableInfo rep -> Context rep)
-> VariableInfo rep -> Context rep
forall a b. (a -> b) -> a -> b
$ Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
forall {k} (rep :: k).
Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
VariableInfo Names
forall a. Monoid a => a
mempty (Context rep -> Int
forall {k} (rep :: k). Context rep -> Int
currentLevel Context rep
ctx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) (Context rep -> [BodyType]
forall {k} (rep :: k). Context rep -> [BodyType]
parents Context rep
ctx') VarType
ThreadID)
          ([(VName, Int)] -> [Context rep])
-> (SegSpace -> [(VName, Int)]) -> SegSpace -> [Context rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\[VName]
segspace_params -> [VName] -> [Int] -> [(VName, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
segspace_params [Int
0 ..])
          -- contextFromNames ctx' Parallel
          ([VName] -> [(VName, Int)])
-> (SegSpace -> [VName]) -> SegSpace -> [(VName, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst
          ([(VName, SubExp)] -> [VName])
-> (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegSpace -> [(VName, SubExp)]
unSegSpace
          (SegSpace -> Context rep) -> SegSpace -> Context rep
forall a b. (a -> b) -> a -> b
$ SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op
   in -- Analyse statements in the SegOp body
      Context rep
-> (VName -> BodyType)
-> [VName]
-> [Stm rep]
-> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep
-> (VName -> BodyType)
-> [VName]
-> [Stm rep]
-> (Context rep, IndexTable rep)
analyseStms Context rep
segspace_context (SegOpName -> BodyType
SegOpName (SegOpName -> BodyType)
-> (VName -> SegOpName) -> VName -> BodyType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegOp lvl rep -> VName -> SegOpName
forall lvl rep. SegOp lvl rep -> VName -> SegOpName
segOpType SegOp lvl rep
op) [VName]
pats ([Stm rep] -> (Context rep, IndexTable rep))
-> (KernelBody rep -> [Stm rep])
-> KernelBody rep
-> (Context rep, IndexTable rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms rep -> [Stm rep])
-> (KernelBody rep -> Stms rep) -> KernelBody rep -> [Stm rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelBody rep -> Stms rep
forall rep. KernelBody rep -> Stms rep
kernelBodyStms (KernelBody rep -> (Context rep, IndexTable rep))
-> KernelBody rep -> (Context rep, IndexTable rep)
forall a b. (a -> b) -> a -> b
$ SegOp lvl rep -> KernelBody rep
forall lvl rep. SegOp lvl rep -> KernelBody rep
segBody SegOp lvl rep
op

analyseSizeOp :: SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSizeOp :: forall {k} (rep :: k).
SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSizeOp SizeOp
op Context rep
ctx [VName]
pats =
  let ctx' :: Context rep
ctx' = case SizeOp
op of
        CmpSizeLe Name
_name SizeClass
_class SubExp
subexp -> [SubExp] -> Context rep
subexprsToContext [SubExp
subexp]
        CalcNumBlocks SubExp
lsubexp Name
_name SubExp
rsubexp -> [SubExp] -> Context rep
subexprsToContext [SubExp
lsubexp, SubExp
rsubexp]
        SizeOp
_ -> Context rep
ctx
      -- Add sizeOp to context
      ctx'' :: Context rep
ctx'' =
        (Context rep -> Context rep -> Context rep)
-> Context rep -> [Context rep] -> Context rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Context rep -> Context rep -> Context rep
forall k (rep :: k). Context rep -> Context rep -> Context rep
extend Context rep
ctx' ([Context rep] -> Context rep) -> [Context rep] -> Context rep
forall a b. (a -> b) -> a -> b
$
          (VName -> Context rep) -> [VName] -> [Context rep]
forall a b. (a -> b) -> [a] -> [b]
map
            (\VName
pat -> VName -> VariableInfo rep -> Context rep
forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
oneContext VName
pat (VariableInfo rep -> Context rep)
-> VariableInfo rep -> Context rep
forall a b. (a -> b) -> a -> b
$ (Context rep -> VariableInfo rep
forall {k} (rep :: k). Context rep -> VariableInfo rep
varInfoZeroDeps Context rep
ctx) {parents_nest = parents ctx'})
            [VName]
pats
   in (Context rep
ctx'', IndexTable rep
forall a. Monoid a => a
mempty)
  where
    subexprsToContext :: [SubExp] -> Context rep
subexprsToContext =
      Context rep -> VariableInfo rep -> [VName] -> Context rep
forall {k} (rep :: k).
Context rep -> VariableInfo rep -> [VName] -> Context rep
contextFromNames Context rep
ctx (Context rep -> VariableInfo rep
forall {k} (rep :: k). Context rep -> VariableInfo rep
varInfoZeroDeps Context rep
ctx)
        ([VName] -> Context rep)
-> ([SubExp] -> [VName]) -> [SubExp] -> Context rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> [VName]) -> [SubExp] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Names -> [VName]
namesToList (Names -> [VName]) -> (SubExp -> Names) -> SubExp -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> Context rep -> SubExp -> Names
forall {k} (rep :: k). [VName] -> Context rep -> SubExp -> Names
analyseSubExp [VName]
pats Context rep
ctx)

-- | Analyse statements in a rep body.
analyseGPUBody :: (Analyse rep) => Body rep -> Context rep -> (Context rep, IndexTable rep)
analyseGPUBody :: forall rep.
Analyse rep =>
Body rep -> Context rep -> (Context rep, IndexTable rep)
analyseGPUBody Body rep
body Context rep
ctx =
  Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
forall rep.
Analyse rep =>
Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
analyseStmsPrimitive Context rep
ctx ([Stm rep] -> (Context rep, IndexTable rep))
-> [Stm rep] -> (Context rep, IndexTable rep)
forall a b. (a -> b) -> a -> b
$ Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms rep -> [Stm rep]) -> Stms rep -> [Stm rep]
forall a b. (a -> b) -> a -> b
$ Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
body

analyseOtherOp :: Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseOtherOp :: forall {k} (rep :: k).
Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseOtherOp Context rep
ctx [VName]
_ = (Context rep
ctx, IndexTable rep
forall a. Monoid a => a
mempty)

-- | Returns an intmap of names, to be used as dependencies in construction of
-- VariableInfos.
analyseSubExp :: [VName] -> Context rep -> SubExp -> Names
analyseSubExp :: forall {k} (rep :: k). [VName] -> Context rep -> SubExp -> Names
analyseSubExp [VName]
_ Context rep
_ (Constant PrimValue
_) = Names
forall a. Monoid a => a
mempty
analyseSubExp [VName]
_ Context rep
_ (Var VName
v) = VName -> Names
oneName VName
v

-- | Reduce a DimFix into its set of dependencies
consolidate :: Context rep -> SubExp -> DimAccess rep
consolidate :: forall {k} (rep :: k). Context rep -> SubExp -> DimAccess rep
consolidate Context rep
_ (Constant PrimValue
_) = DimAccess rep
forall a. Monoid a => a
mempty
consolidate Context rep
ctx (Var VName
v) = Map VName Dependency -> Maybe VName -> DimAccess rep
forall {k} (rep :: k).
Map VName Dependency -> Maybe VName -> DimAccess rep
DimAccess (Context rep -> VName -> Map VName Dependency
forall {k} (rep :: k). Context rep -> VName -> Map VName Dependency
reduceDependencies Context rep
ctx VName
v) (VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v)

-- | Recursively lookup vnames until vars with no deps are reached.
reduceDependencies :: Context rep -> VName -> M.Map VName Dependency
reduceDependencies :: forall {k} (rep :: k). Context rep -> VName -> Map VName Dependency
reduceDependencies Context rep
ctx VName
v =
  case VName -> Map VName (VariableInfo rep) -> Maybe (VariableInfo rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Context rep -> Map VName (VariableInfo rep)
forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
assignments Context rep
ctx) of
    Maybe (VariableInfo rep)
Nothing -> Map VName Dependency
forall a. Monoid a => a
mempty -- Means a global.
    Just (VariableInfo Names
deps Int
lvl [BodyType]
_parents VarType
t) ->
      -- We detect whether it is a threadID or loop counter by checking
      -- whether or not it has any dependencies
      case VarType
t of
        VarType
ThreadID -> [(VName, Dependency)] -> Map VName Dependency
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName
v, Int -> VarType -> Dependency
Dependency Int
lvl VarType
t)]
        VarType
LoopVar -> [(VName, Dependency)] -> Map VName Dependency
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName
v, Int -> VarType -> Dependency
Dependency Int
lvl VarType
t)]
        VarType
Variable -> [Map VName Dependency] -> Map VName Dependency
forall a. Monoid a => [a] -> a
mconcat ([Map VName Dependency] -> Map VName Dependency)
-> [Map VName Dependency] -> Map VName Dependency
forall a b. (a -> b) -> a -> b
$ (VName -> Map VName Dependency)
-> [VName] -> [Map VName Dependency]
forall a b. (a -> b) -> [a] -> [b]
map (Context rep -> VName -> Map VName Dependency
forall {k} (rep :: k). Context rep -> VName -> Map VName Dependency
reduceDependencies Context rep
ctx) ([VName] -> [Map VName Dependency])
-> [VName] -> [Map VName Dependency]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
deps
        VarType
ConstType -> Map VName Dependency
forall a. Monoid a => a
mempty

-- Misc functions

-- Instances for AST types that we actually support
instance Analyse GPU where
  analyseOp :: Op GPU -> Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
analyseOp Op GPU
gpu_op
    | (SegOp SegOp SegLevel GPU
op) <- Op GPU
gpu_op = SegOp SegLevel GPU
-> Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
forall rep lvl.
Analyse rep =>
SegOp lvl rep
-> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSegOp SegOp SegLevel GPU
op
    | (SizeOp SizeOp
op) <- Op GPU
gpu_op = SizeOp -> Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
forall {k} (rep :: k).
SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSizeOp SizeOp
op
    | (GPUBody [Type]
_ Body GPU
body) <- Op GPU
gpu_op = (Context GPU, IndexTable GPU)
-> [VName] -> (Context GPU, IndexTable GPU)
forall a. a -> [VName] -> a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Context GPU, IndexTable GPU)
 -> [VName] -> (Context GPU, IndexTable GPU))
-> (Context GPU -> (Context GPU, IndexTable GPU))
-> Context GPU
-> [VName]
-> (Context GPU, IndexTable GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body GPU -> Context GPU -> (Context GPU, IndexTable GPU)
forall rep.
Analyse rep =>
Body rep -> Context rep -> (Context rep, IndexTable rep)
analyseGPUBody Body GPU
body
    | (Futhark.IR.GPU.OtherOp SOAC GPU
_) <- Op GPU
gpu_op = Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
forall {k} (rep :: k).
Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseOtherOp

instance Analyse MC where
  analyseOp :: Op MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
analyseOp Op MC
mc_op
    | ParOp Maybe (SegOp () MC)
Nothing SegOp () MC
seq_segop <- Op MC
mc_op = SegOp () MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
forall rep lvl.
Analyse rep =>
SegOp lvl rep
-> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSegOp SegOp () MC
seq_segop
    | ParOp (Just SegOp () MC
segop) SegOp () MC
seq_segop <- Op MC
mc_op = \Context MC
ctx [VName]
name -> do
        let (Context MC
ctx', IndexTable MC
res') = SegOp () MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
forall rep lvl.
Analyse rep =>
SegOp lvl rep
-> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSegOp SegOp () MC
segop Context MC
ctx [VName]
name
        let (Context MC
ctx'', IndexTable MC
res'') = SegOp () MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
forall rep lvl.
Analyse rep =>
SegOp lvl rep
-> Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseSegOp SegOp () MC
seq_segop Context MC
ctx' [VName]
name
        (Context MC
ctx'', IndexTable MC -> IndexTable MC -> IndexTable MC
forall {k} (rep :: k).
IndexTable rep -> IndexTable rep -> IndexTable rep
unionIndexTables IndexTable MC
res' IndexTable MC
res'')
    | Futhark.IR.MC.OtherOp SOAC MC
_ <- Op MC
mc_op = Context MC -> [VName] -> (Context MC, IndexTable MC)
forall {k} (rep :: k).
Context rep -> [VName] -> (Context rep, IndexTable rep)
analyseOtherOp

-- Unfortunately we need these instances, even though they may never appear.
instance Analyse GPUMem where
  analyseOp :: Op GPUMem
-> Context GPUMem -> [VName] -> (Context GPUMem, IndexTable GPUMem)
analyseOp Op GPUMem
_ = [Char]
-> Context GPUMem -> [VName] -> (Context GPUMem, IndexTable GPUMem)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> Context GPUMem
 -> [VName]
 -> (Context GPUMem, IndexTable GPUMem))
-> [Char]
-> Context GPUMem
-> [VName]
-> (Context GPUMem, IndexTable GPUMem)
forall a b. (a -> b) -> a -> b
$ ShowS
notImplementedYet [Char]
"GPUMem"

instance Analyse MCMem where
  analyseOp :: Op MCMem
-> Context MCMem -> [VName] -> (Context MCMem, IndexTable MCMem)
analyseOp Op MCMem
_ = [Char]
-> Context MCMem -> [VName] -> (Context MCMem, IndexTable MCMem)
forall a. HasCallStack => [Char] -> a
error [Char]
"Unexpected?"

instance Analyse Seq where
  analyseOp :: Op Seq -> Context Seq -> [VName] -> (Context Seq, IndexTable Seq)
analyseOp Op Seq
_ = [Char] -> Context Seq -> [VName] -> (Context Seq, IndexTable Seq)
forall a. HasCallStack => [Char] -> a
error ([Char] -> Context Seq -> [VName] -> (Context Seq, IndexTable Seq))
-> [Char]
-> Context Seq
-> [VName]
-> (Context Seq, IndexTable Seq)
forall a b. (a -> b) -> a -> b
$ ShowS
notImplementedYet [Char]
"Seq"

instance Analyse SeqMem where
  analyseOp :: Op SeqMem
-> Context SeqMem -> [VName] -> (Context SeqMem, IndexTable SeqMem)
analyseOp Op SeqMem
_ = [Char]
-> Context SeqMem -> [VName] -> (Context SeqMem, IndexTable SeqMem)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> Context SeqMem
 -> [VName]
 -> (Context SeqMem, IndexTable SeqMem))
-> [Char]
-> Context SeqMem
-> [VName]
-> (Context SeqMem, IndexTable SeqMem)
forall a b. (a -> b) -> a -> b
$ ShowS
notImplementedYet [Char]
"SeqMem"

instance Analyse SOACS where
  analyseOp :: Op SOACS
-> Context SOACS -> [VName] -> (Context SOACS, IndexTable SOACS)
analyseOp Op SOACS
_ = [Char]
-> Context SOACS -> [VName] -> (Context SOACS, IndexTable SOACS)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> Context SOACS -> [VName] -> (Context SOACS, IndexTable SOACS))
-> [Char]
-> Context SOACS
-> [VName]
-> (Context SOACS, IndexTable SOACS)
forall a b. (a -> b) -> a -> b
$ ShowS
notImplementedYet [Char]
"SOACS"

notImplementedYet :: String -> String
notImplementedYet :: ShowS
notImplementedYet [Char]
s = [Char]
"Access pattern analysis for the " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
s [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" backend is not implemented."

instance Pretty (IndexTable rep) where
  pretty :: forall ann. IndexTable rep -> Doc ann
pretty = [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
stack ([Doc ann] -> Doc ann)
-> (IndexTable rep -> [Doc ann]) -> IndexTable rep -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((SegOpName, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann)
-> [(SegOpName, Map ArrayName (Map VName [DimAccess rep]))]
-> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (SegOpName, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann
forall {a} {ann}.
Pretty a =>
(a, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann
f ([(SegOpName, Map ArrayName (Map VName [DimAccess rep]))]
 -> [Doc ann])
-> (IndexTable rep
    -> [(SegOpName, Map ArrayName (Map VName [DimAccess rep]))])
-> IndexTable rep
-> [Doc ann]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IndexTable rep
-> [(SegOpName, Map ArrayName (Map VName [DimAccess rep]))]
forall k a. Map k a -> [(k, a)]
M.toList :: IndexTable rep -> Doc ann
    where
      f :: (a, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann
f (a
segop, Map ArrayName (Map VName [DimAccess rep])
arrNameToIdxExprMap) = a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
segop Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Map ArrayName (Map VName [DimAccess rep]) -> Doc ann
forall {a}. Map ArrayName (Map VName [DimAccess rep]) -> Doc a
g Map ArrayName (Map VName [DimAccess rep])
arrNameToIdxExprMap

      g :: Map ArrayName (Map VName [DimAccess rep]) -> Doc a
g Map ArrayName (Map VName [DimAccess rep])
maps = Doc a
forall ann. Doc ann
lbrace Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc a -> Doc a
forall ann. Int -> Doc ann -> Doc ann
indent Int
4 ([(ArrayName, Map VName [DimAccess rep])] -> Doc a
forall ann. [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
mapprintArray ([(ArrayName, Map VName [DimAccess rep])] -> Doc a)
-> [(ArrayName, Map VName [DimAccess rep])] -> Doc a
forall a b. (a -> b) -> a -> b
$ Map ArrayName (Map VName [DimAccess rep])
-> [(ArrayName, Map VName [DimAccess rep])]
forall k a. Map k a -> [(k, a)]
M.toList Map ArrayName (Map VName [DimAccess rep])
maps) Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc a
forall ann. Doc ann
rbrace

      mapprintArray :: [(ArrayName, M.Map IndexExprName [DimAccess rep])] -> Doc ann
      mapprintArray :: forall ann. [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
mapprintArray [] = Doc ann
""
      mapprintArray [(ArrayName, Map VName [DimAccess rep])
m] = (ArrayName, Map VName [DimAccess rep]) -> Doc ann
forall ann. (ArrayName, Map VName [DimAccess rep]) -> Doc ann
printArrayMap (ArrayName, Map VName [DimAccess rep])
m
      mapprintArray ((ArrayName, Map VName [DimAccess rep])
m : [(ArrayName, Map VName [DimAccess rep])]
mm) = (ArrayName, Map VName [DimAccess rep]) -> Doc ann
forall ann. (ArrayName, Map VName [DimAccess rep]) -> Doc ann
printArrayMap (ArrayName, Map VName [DimAccess rep])
m Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
</> [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
forall ann. [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
mapprintArray [(ArrayName, Map VName [DimAccess rep])]
mm

      printArrayMap :: (ArrayName, M.Map IndexExprName [DimAccess rep]) -> Doc ann
      printArrayMap :: forall ann. (ArrayName, Map VName [DimAccess rep]) -> Doc ann
printArrayMap ((VName
name, [BodyType]
_, [Int]
layout), Map VName [DimAccess rep]
maps) =
        Doc ann
"(arr)"
          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
name
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Int] -> Doc ann
forall ann. [Int] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [Int]
layout
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
lbrace
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
4 ([(VName, [DimAccess rep])] -> Doc ann
forall ann. [(VName, [DimAccess rep])] -> Doc ann
mapprintIdxExpr (Map VName [DimAccess rep] -> [(VName, [DimAccess rep])]
forall k a. Map k a -> [(k, a)]
M.toList Map VName [DimAccess rep]
maps))
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ann
forall ann. Doc ann
rbrace

      mapprintIdxExpr :: [(IndexExprName, [DimAccess rep])] -> Doc ann
      mapprintIdxExpr :: forall ann. [(VName, [DimAccess rep])] -> Doc ann
mapprintIdxExpr [] = Doc ann
""
      mapprintIdxExpr [(VName, [DimAccess rep])
m] = (VName, [DimAccess rep]) -> Doc ann
forall {a} {ann}. Pretty a => (a, [DimAccess rep]) -> Doc ann
printIdxExpMap (VName, [DimAccess rep])
m
      mapprintIdxExpr ((VName, [DimAccess rep])
m : [(VName, [DimAccess rep])]
mm) = (VName, [DimAccess rep]) -> Doc ann
forall {a} {ann}. Pretty a => (a, [DimAccess rep]) -> Doc ann
printIdxExpMap (VName, [DimAccess rep])
m Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
</> [(VName, [DimAccess rep])] -> Doc ann
forall ann. [(VName, [DimAccess rep])] -> Doc ann
mapprintIdxExpr [(VName, [DimAccess rep])]
mm

      printIdxExpMap :: (a, [DimAccess rep]) -> Doc ann
printIdxExpMap (a
name, [DimAccess rep]
mems) = Doc ann
"(idx)" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
name 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
</> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
4 ([DimAccess rep] -> Doc ann
forall ann. [DimAccess rep] -> Doc ann
printDimAccess [DimAccess rep]
mems)

      printDimAccess :: [DimAccess rep] -> Doc ann
      printDimAccess :: forall ann. [DimAccess rep] -> Doc ann
printDimAccess [DimAccess rep]
dim_accesses = [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
stack ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ (Int -> DimAccess rep -> Doc ann)
-> [Int] -> [DimAccess rep] -> [Doc ann]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (((Int, DimAccess rep) -> Doc ann)
-> Int -> DimAccess rep -> Doc ann
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (Int, DimAccess rep) -> Doc ann
forall ann. (Int, DimAccess rep) -> Doc ann
printDim) [Int
0 ..] [DimAccess rep]
dim_accesses

      printDim :: (Int, DimAccess rep) -> Doc ann
      printDim :: forall ann. (Int, DimAccess rep) -> Doc ann
printDim (Int
i, DimAccess rep
m) = Int -> Doc ann
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Int
i 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
<+> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
0 (DimAccess rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. DimAccess rep -> Doc ann
pretty DimAccess rep
m)

instance Pretty (DimAccess rep) where
  pretty :: forall ann. DimAccess rep -> Doc ann
pretty DimAccess rep
dim_access =
    -- Instead of using `brackets $` we manually enclose with `[`s, to add
    -- spacing between the enclosed elements
    if case DimAccess rep -> Maybe VName
forall {k} (rep :: k). DimAccess rep -> Maybe VName
originalVar DimAccess rep
dim_access of
      Maybe VName
Nothing -> Bool
True
      Just VName
n ->
        Map VName Dependency -> Int
forall a. Map VName a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (DimAccess rep -> Map VName Dependency
forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies DimAccess rep
dim_access) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& VName
n VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== [VName] -> VName
forall a. HasCallStack => [a] -> a
head (((VName, Dependency) -> VName) -> [(VName, Dependency)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Dependency) -> VName
forall a b. (a, b) -> a
fst ([(VName, Dependency)] -> [VName])
-> [(VName, Dependency)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Map VName Dependency -> [(VName, Dependency)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName Dependency -> [(VName, Dependency)])
-> Map VName Dependency -> [(VName, Dependency)]
forall a b. (a -> b) -> a -> b
$ DimAccess rep -> Map VName Dependency
forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies DimAccess rep
dim_access)
        -- Only print the original name if it is different from the first (and single) dependency
      then
        Doc ann
"dependencies"
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
equals
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align (Map VName Dependency -> Doc ann
forall {ann}. Map VName Dependency -> Doc ann
prettyDeps (Map VName Dependency -> Doc ann)
-> Map VName Dependency -> Doc ann
forall a b. (a -> b) -> a -> b
$ DimAccess rep -> Map VName Dependency
forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies DimAccess rep
dim_access)
      else
        Doc ann
"dependencies"
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
equals
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Maybe VName -> Doc ann
forall ann. Maybe VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (DimAccess rep -> Maybe VName
forall {k} (rep :: k). DimAccess rep -> Maybe VName
originalVar DimAccess rep
dim_access)
          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 -> Doc ann
forall ann. Doc ann -> Doc ann
align (Map VName Dependency -> Doc ann
forall {ann}. Map VName Dependency -> Doc ann
prettyDeps (Map VName Dependency -> Doc ann)
-> Map VName Dependency -> Doc ann
forall a b. (a -> b) -> a -> b
$ DimAccess rep -> Map VName Dependency
forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
dependencies DimAccess rep
dim_access)
    where
      prettyDeps :: Map VName Dependency -> Doc ann
prettyDeps = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
braces (Doc ann -> Doc ann)
-> (Map VName Dependency -> Doc ann)
-> Map VName Dependency
-> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
commasep ([Doc ann] -> Doc ann)
-> (Map VName Dependency -> [Doc ann])
-> Map VName Dependency
-> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, Dependency) -> Doc ann)
-> [(VName, Dependency)] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Dependency) -> Doc ann
forall {a} {ann}. Pretty a => (a, Dependency) -> Doc ann
printPair ([(VName, Dependency)] -> [Doc ann])
-> (Map VName Dependency -> [(VName, Dependency)])
-> Map VName Dependency
-> [Doc ann]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Dependency -> [(VName, Dependency)]
forall k a. Map k a -> [(k, a)]
M.toList
      printPair :: (a, Dependency) -> Doc ann
printPair (a
name, Dependency Int
lvl VarType
vtype) = a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Int -> Doc ann
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Int
lvl Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> VarType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VarType -> Doc ann
pretty VarType
vtype

instance Pretty SegOpName where
  pretty :: forall ann. SegOpName -> Doc ann
pretty (SegmentedMap VName
name) = Doc ann
"(segmap)" 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
name
  pretty (SegmentedRed VName
name) = Doc ann
"(segred)" 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
name
  pretty (SegmentedScan VName
name) = Doc ann
"(segscan)" 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
name
  pretty (SegmentedHist VName
name) = Doc ann
"(seghist)" 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
name

instance Pretty BodyType where
  pretty :: forall ann. BodyType -> Doc ann
pretty (SegOpName (SegmentedMap VName
name)) = VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"segmap"
  pretty (SegOpName (SegmentedRed VName
name)) = VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"segred"
  pretty (SegOpName (SegmentedScan VName
name)) = VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"segscan"
  pretty (SegOpName (SegmentedHist VName
name)) = VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"seghist"
  pretty (LoopBodyName VName
name) = VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"loop"
  pretty (CondBodyName VName
name) = VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"cond"

instance Pretty VarType where
  pretty :: forall ann. VarType -> Doc ann
pretty VarType
ConstType = Doc ann
"const"
  pretty VarType
Variable = Doc ann
"var"
  pretty VarType
ThreadID = Doc ann
"tid"
  pretty VarType
LoopVar = Doc ann
"iter"