{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Segmented operations.  These correspond to perfect @map@ nests on
-- top of /something/, except that the @map@s are conceptually only
-- over @iota@s (so there will be explicit indexing inside them).
module Futhark.IR.SegOp
  ( SegOp (..),
    segLevel,
    segBody,
    segSpace,
    typeCheckSegOp,
    SegSpace (..),
    scopeOfSegSpace,
    segSpaceDims,

    -- * Details
    HistOp (..),
    histType,
    splitHistResults,
    SegBinOp (..),
    segBinOpResults,
    segBinOpChunks,
    KernelBody,
    ResultManifest (..),
    KernelResult (..),
    kernelResultCerts,
    kernelResultSubExp,
    SegPostOp (..),

    -- ** Generic traversal
    SegOpMapper (..),
    identitySegOpMapper,
    mapSegOpM,
    traverseSegOpStms,

    -- * Simplification
    simplifySegOp,
    HasSegOp (..),
    segOpRules,

    -- * Memory
    segOpReturns,
  )
where

import Control.Category
import Control.Monad
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State.Strict
import Control.Monad.Writer
import Data.Bifunctor (first)
import Data.Bitraversable
import Data.List
  ( elemIndex,
    groupBy,
    intersperse,
    isPrefixOf,
    partition,
    unzip4,
    zip4,
  )
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.Alias qualified as Alias
import Futhark.Analysis.DataDependencies
import Futhark.Analysis.Metrics
import Futhark.Analysis.PrimExp.Convert
import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.Analysis.UsageTable qualified as UT
import Futhark.IR
import Futhark.IR.Aliases
  ( Aliases,
    CanBeAliased (..),
  )
import Futhark.IR.Mem
import Futhark.IR.Prop.Aliases
import Futhark.IR.TypeCheck qualified as TC
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Optimise.Simplify.Rep
import Futhark.Optimise.Simplify.Rule
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util (chunks, maybeNth)
import Futhark.Util.Pretty
  ( Doc,
    apply,
    hsep,
    parens,
    ppTuple',
    pretty,
    (<+>),
    (</>),
  )
import Futhark.Util.Pretty qualified as PP
import Prelude hiding (id, (.))

-- | An operator for 'SegHist'.
data HistOp rep = HistOp
  { forall rep. HistOp rep -> Shape
histShape :: Shape,
    forall rep. HistOp rep -> SubExp
histRaceFactor :: SubExp,
    forall rep. HistOp rep -> [VName]
histDest :: [VName],
    forall rep. HistOp rep -> [SubExp]
histNeutral :: [SubExp],
    -- | In case this operator is semantically a vectorised
    -- operator (corresponding to a perfect map nest in the
    -- SOACS representation), these are the logical
    -- "dimensions".  This is used to generate more efficient
    -- code.
    forall rep. HistOp rep -> Shape
histOpShape :: Shape,
    forall rep. HistOp rep -> Lambda rep
histOp :: Lambda rep
  }
  deriving (HistOp rep -> HistOp rep -> Bool
(HistOp rep -> HistOp rep -> Bool)
-> (HistOp rep -> HistOp rep -> Bool) -> Eq (HistOp rep)
forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
== :: HistOp rep -> HistOp rep -> Bool
$c/= :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
/= :: HistOp rep -> HistOp rep -> Bool
Eq, Eq (HistOp rep)
Eq (HistOp rep) =>
(HistOp rep -> HistOp rep -> Ordering)
-> (HistOp rep -> HistOp rep -> Bool)
-> (HistOp rep -> HistOp rep -> Bool)
-> (HistOp rep -> HistOp rep -> Bool)
-> (HistOp rep -> HistOp rep -> Bool)
-> (HistOp rep -> HistOp rep -> HistOp rep)
-> (HistOp rep -> HistOp rep -> HistOp rep)
-> Ord (HistOp rep)
HistOp rep -> HistOp rep -> Bool
HistOp rep -> HistOp rep -> Ordering
HistOp rep -> HistOp rep -> HistOp rep
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
forall rep. RepTypes rep => Eq (HistOp rep)
forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Ordering
forall rep. RepTypes rep => HistOp rep -> HistOp rep -> HistOp rep
$ccompare :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Ordering
compare :: HistOp rep -> HistOp rep -> Ordering
$c< :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
< :: HistOp rep -> HistOp rep -> Bool
$c<= :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
<= :: HistOp rep -> HistOp rep -> Bool
$c> :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
> :: HistOp rep -> HistOp rep -> Bool
$c>= :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
>= :: HistOp rep -> HistOp rep -> Bool
$cmax :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> HistOp rep
max :: HistOp rep -> HistOp rep -> HistOp rep
$cmin :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> HistOp rep
min :: HistOp rep -> HistOp rep -> HistOp rep
Ord, Int -> HistOp rep -> ShowS
[HistOp rep] -> ShowS
HistOp rep -> String
(Int -> HistOp rep -> ShowS)
-> (HistOp rep -> String)
-> ([HistOp rep] -> ShowS)
-> Show (HistOp rep)
forall rep. RepTypes rep => Int -> HistOp rep -> ShowS
forall rep. RepTypes rep => [HistOp rep] -> ShowS
forall rep. RepTypes rep => HistOp rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> HistOp rep -> ShowS
showsPrec :: Int -> HistOp rep -> ShowS
$cshow :: forall rep. RepTypes rep => HistOp rep -> String
show :: HistOp rep -> String
$cshowList :: forall rep. RepTypes rep => [HistOp rep] -> ShowS
showList :: [HistOp rep] -> ShowS
Show)

-- | The type of a histogram produced by a 'HistOp'.  This can be
-- different from the type of the 'histDest's in case we are
-- dealing with a segmented histogram.
histType :: HistOp rep -> [Type]
histType :: forall rep. HistOp rep -> [TypeBase Shape NoMode]
histType HistOp rep
op =
  (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
`arrayOfShape` (HistOp rep -> Shape
forall rep. HistOp rep -> Shape
histShape HistOp rep
op Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> HistOp rep -> Shape
forall rep. HistOp rep -> Shape
histOpShape HistOp rep
op)) ([TypeBase Shape NoMode] -> [TypeBase Shape NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$
    GLambda rep (TypeBase Shape NoMode) -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (GLambda rep (TypeBase Shape NoMode) -> [TypeBase Shape NoMode])
-> GLambda rep (TypeBase Shape NoMode) -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$
      HistOp rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. HistOp rep -> Lambda rep
histOp HistOp rep
op

-- | Split reduction results returned by a 'KernelBody' into those
-- that correspond to indexes for the 'HistOp's, and those that
-- correspond to value.
splitHistResults :: [HistOp rep] -> [SubExp] -> [([SubExp], [SubExp])]
splitHistResults :: forall rep. [HistOp rep] -> [SubExp] -> [([SubExp], [SubExp])]
splitHistResults [HistOp rep]
ops [SubExp]
res =
  let ranks :: [Int]
ranks = (HistOp rep -> Int) -> [HistOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> (HistOp rep -> Shape) -> HistOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. HistOp rep -> Shape
forall rep. HistOp rep -> Shape
histShape) [HistOp rep]
ops
      ([SubExp]
idxs, [SubExp]
vals) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
ranks) [SubExp]
res
   in [[SubExp]] -> [[SubExp]] -> [([SubExp], [SubExp])]
forall a b. [a] -> [b] -> [(a, b)]
zip
        ([Int] -> [SubExp] -> [[SubExp]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
ranks [SubExp]
idxs)
        ([Int] -> [SubExp] -> [[SubExp]]
forall a. [Int] -> [a] -> [[a]]
chunks ((HistOp rep -> Int) -> [HistOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([VName] -> Int) -> (HistOp rep -> [VName]) -> HistOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. HistOp rep -> [VName]
forall rep. HistOp rep -> [VName]
histDest) [HistOp rep]
ops) [SubExp]
vals)

-- | An operator for 'SegScan' and 'SegRed'.
data SegBinOp rep = SegBinOp
  { forall rep. SegBinOp rep -> Commutativity
segBinOpComm :: Commutativity,
    forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda :: Lambda rep,
    forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral :: [SubExp],
    -- | In case this operator is semantically a vectorised
    -- operator (corresponding to a perfect map nest in the
    -- SOACS representation), these are the logical
    -- "dimensions".  This is used to generate more efficient
    -- code.
    forall rep. SegBinOp rep -> Shape
segBinOpShape :: Shape
  }
  deriving (SegBinOp rep -> SegBinOp rep -> Bool
(SegBinOp rep -> SegBinOp rep -> Bool)
-> (SegBinOp rep -> SegBinOp rep -> Bool) -> Eq (SegBinOp rep)
forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
== :: SegBinOp rep -> SegBinOp rep -> Bool
$c/= :: forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
/= :: SegBinOp rep -> SegBinOp rep -> Bool
Eq, Eq (SegBinOp rep)
Eq (SegBinOp rep) =>
(SegBinOp rep -> SegBinOp rep -> Ordering)
-> (SegBinOp rep -> SegBinOp rep -> Bool)
-> (SegBinOp rep -> SegBinOp rep -> Bool)
-> (SegBinOp rep -> SegBinOp rep -> Bool)
-> (SegBinOp rep -> SegBinOp rep -> Bool)
-> (SegBinOp rep -> SegBinOp rep -> SegBinOp rep)
-> (SegBinOp rep -> SegBinOp rep -> SegBinOp rep)
-> Ord (SegBinOp rep)
SegBinOp rep -> SegBinOp rep -> Bool
SegBinOp rep -> SegBinOp rep -> Ordering
SegBinOp rep -> SegBinOp rep -> SegBinOp rep
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
forall rep. RepTypes rep => Eq (SegBinOp rep)
forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
forall rep.
RepTypes rep =>
SegBinOp rep -> SegBinOp rep -> Ordering
forall rep.
RepTypes rep =>
SegBinOp rep -> SegBinOp rep -> SegBinOp rep
$ccompare :: forall rep.
RepTypes rep =>
SegBinOp rep -> SegBinOp rep -> Ordering
compare :: SegBinOp rep -> SegBinOp rep -> Ordering
$c< :: forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
< :: SegBinOp rep -> SegBinOp rep -> Bool
$c<= :: forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
<= :: SegBinOp rep -> SegBinOp rep -> Bool
$c> :: forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
> :: SegBinOp rep -> SegBinOp rep -> Bool
$c>= :: forall rep. RepTypes rep => SegBinOp rep -> SegBinOp rep -> Bool
>= :: SegBinOp rep -> SegBinOp rep -> Bool
$cmax :: forall rep.
RepTypes rep =>
SegBinOp rep -> SegBinOp rep -> SegBinOp rep
max :: SegBinOp rep -> SegBinOp rep -> SegBinOp rep
$cmin :: forall rep.
RepTypes rep =>
SegBinOp rep -> SegBinOp rep -> SegBinOp rep
min :: SegBinOp rep -> SegBinOp rep -> SegBinOp rep
Ord, Int -> SegBinOp rep -> ShowS
[SegBinOp rep] -> ShowS
SegBinOp rep -> String
(Int -> SegBinOp rep -> ShowS)
-> (SegBinOp rep -> String)
-> ([SegBinOp rep] -> ShowS)
-> Show (SegBinOp rep)
forall rep. RepTypes rep => Int -> SegBinOp rep -> ShowS
forall rep. RepTypes rep => [SegBinOp rep] -> ShowS
forall rep. RepTypes rep => SegBinOp rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> SegBinOp rep -> ShowS
showsPrec :: Int -> SegBinOp rep -> ShowS
$cshow :: forall rep. RepTypes rep => SegBinOp rep -> String
show :: SegBinOp rep -> String
$cshowList :: forall rep. RepTypes rep => [SegBinOp rep] -> ShowS
showList :: [SegBinOp rep] -> ShowS
Show)

newtype SegPostOp rep = SegPostOp
  { -- | The function applied the the result of some SegBinOp.
    forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda :: Lambda rep
  }
  deriving (SegPostOp rep -> SegPostOp rep -> Bool
(SegPostOp rep -> SegPostOp rep -> Bool)
-> (SegPostOp rep -> SegPostOp rep -> Bool) -> Eq (SegPostOp rep)
forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
== :: SegPostOp rep -> SegPostOp rep -> Bool
$c/= :: forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
/= :: SegPostOp rep -> SegPostOp rep -> Bool
Eq, Eq (SegPostOp rep)
Eq (SegPostOp rep) =>
(SegPostOp rep -> SegPostOp rep -> Ordering)
-> (SegPostOp rep -> SegPostOp rep -> Bool)
-> (SegPostOp rep -> SegPostOp rep -> Bool)
-> (SegPostOp rep -> SegPostOp rep -> Bool)
-> (SegPostOp rep -> SegPostOp rep -> Bool)
-> (SegPostOp rep -> SegPostOp rep -> SegPostOp rep)
-> (SegPostOp rep -> SegPostOp rep -> SegPostOp rep)
-> Ord (SegPostOp rep)
SegPostOp rep -> SegPostOp rep -> Bool
SegPostOp rep -> SegPostOp rep -> Ordering
SegPostOp rep -> SegPostOp rep -> SegPostOp rep
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
forall rep. RepTypes rep => Eq (SegPostOp rep)
forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
forall rep.
RepTypes rep =>
SegPostOp rep -> SegPostOp rep -> Ordering
forall rep.
RepTypes rep =>
SegPostOp rep -> SegPostOp rep -> SegPostOp rep
$ccompare :: forall rep.
RepTypes rep =>
SegPostOp rep -> SegPostOp rep -> Ordering
compare :: SegPostOp rep -> SegPostOp rep -> Ordering
$c< :: forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
< :: SegPostOp rep -> SegPostOp rep -> Bool
$c<= :: forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
<= :: SegPostOp rep -> SegPostOp rep -> Bool
$c> :: forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
> :: SegPostOp rep -> SegPostOp rep -> Bool
$c>= :: forall rep. RepTypes rep => SegPostOp rep -> SegPostOp rep -> Bool
>= :: SegPostOp rep -> SegPostOp rep -> Bool
$cmax :: forall rep.
RepTypes rep =>
SegPostOp rep -> SegPostOp rep -> SegPostOp rep
max :: SegPostOp rep -> SegPostOp rep -> SegPostOp rep
$cmin :: forall rep.
RepTypes rep =>
SegPostOp rep -> SegPostOp rep -> SegPostOp rep
min :: SegPostOp rep -> SegPostOp rep -> SegPostOp rep
Ord, Int -> SegPostOp rep -> ShowS
[SegPostOp rep] -> ShowS
SegPostOp rep -> String
(Int -> SegPostOp rep -> ShowS)
-> (SegPostOp rep -> String)
-> ([SegPostOp rep] -> ShowS)
-> Show (SegPostOp rep)
forall rep. RepTypes rep => Int -> SegPostOp rep -> ShowS
forall rep. RepTypes rep => [SegPostOp rep] -> ShowS
forall rep. RepTypes rep => SegPostOp rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> SegPostOp rep -> ShowS
showsPrec :: Int -> SegPostOp rep -> ShowS
$cshow :: forall rep. RepTypes rep => SegPostOp rep -> String
show :: SegPostOp rep -> String
$cshowList :: forall rep. RepTypes rep => [SegPostOp rep] -> ShowS
showList :: [SegPostOp rep] -> ShowS
Show)

-- | How many reduction results are produced by these 'SegBinOp's?
segBinOpResults :: [SegBinOp rep] -> Int
segBinOpResults :: forall rep. [SegBinOp rep] -> Int
segBinOpResults = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int)
-> ([SegBinOp rep] -> [Int]) -> [SegBinOp rep] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (SegBinOp rep -> Int) -> [SegBinOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int)
-> (SegBinOp rep -> [SubExp]) -> SegBinOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral)

-- | Split some list into chunks equal to the number of values
-- returned by each 'SegBinOp'
segBinOpChunks :: [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks :: forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks = [Int] -> [a] -> [[a]]
forall a. [Int] -> [a] -> [[a]]
chunks ([Int] -> [a] -> [[a]])
-> ([SegBinOp rep] -> [Int]) -> [SegBinOp rep] -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (SegBinOp rep -> Int) -> [SegBinOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int)
-> (SegBinOp rep -> [SubExp]) -> SegBinOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral)

-- | The body of a 'SegOp'.
type KernelBody rep = GBody rep KernelResult

-- | Metadata about whether there is a subtle point to this
-- 'KernelResult'.  This is used to protect things like tiling, which
-- might otherwise be removed by the simplifier because they're
-- semantically redundant.  This has no semantic effect and can be
-- ignored at code generation.
data ResultManifest
  = -- | Don't simplify this one!
    ResultNoSimplify
  | -- | Go nuts.
    ResultMaySimplify
  | -- | The results produced are only used within the
    -- same physical thread later on, and can thus be
    -- kept in registers.
    ResultPrivate
  deriving (ResultManifest -> ResultManifest -> Bool
(ResultManifest -> ResultManifest -> Bool)
-> (ResultManifest -> ResultManifest -> Bool) -> Eq ResultManifest
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ResultManifest -> ResultManifest -> Bool
== :: ResultManifest -> ResultManifest -> Bool
$c/= :: ResultManifest -> ResultManifest -> Bool
/= :: ResultManifest -> ResultManifest -> Bool
Eq, Int -> ResultManifest -> ShowS
[ResultManifest] -> ShowS
ResultManifest -> String
(Int -> ResultManifest -> ShowS)
-> (ResultManifest -> String)
-> ([ResultManifest] -> ShowS)
-> Show ResultManifest
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ResultManifest -> ShowS
showsPrec :: Int -> ResultManifest -> ShowS
$cshow :: ResultManifest -> String
show :: ResultManifest -> String
$cshowList :: [ResultManifest] -> ShowS
showList :: [ResultManifest] -> ShowS
Show, Eq ResultManifest
Eq ResultManifest =>
(ResultManifest -> ResultManifest -> Ordering)
-> (ResultManifest -> ResultManifest -> Bool)
-> (ResultManifest -> ResultManifest -> Bool)
-> (ResultManifest -> ResultManifest -> Bool)
-> (ResultManifest -> ResultManifest -> Bool)
-> (ResultManifest -> ResultManifest -> ResultManifest)
-> (ResultManifest -> ResultManifest -> ResultManifest)
-> Ord ResultManifest
ResultManifest -> ResultManifest -> Bool
ResultManifest -> ResultManifest -> Ordering
ResultManifest -> ResultManifest -> ResultManifest
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 :: ResultManifest -> ResultManifest -> Ordering
compare :: ResultManifest -> ResultManifest -> Ordering
$c< :: ResultManifest -> ResultManifest -> Bool
< :: ResultManifest -> ResultManifest -> Bool
$c<= :: ResultManifest -> ResultManifest -> Bool
<= :: ResultManifest -> ResultManifest -> Bool
$c> :: ResultManifest -> ResultManifest -> Bool
> :: ResultManifest -> ResultManifest -> Bool
$c>= :: ResultManifest -> ResultManifest -> Bool
>= :: ResultManifest -> ResultManifest -> Bool
$cmax :: ResultManifest -> ResultManifest -> ResultManifest
max :: ResultManifest -> ResultManifest -> ResultManifest
$cmin :: ResultManifest -> ResultManifest -> ResultManifest
min :: ResultManifest -> ResultManifest -> ResultManifest
Ord)

-- | A 'KernelBody' does not return an ordinary 'Result'.  Instead, it
-- returns a list of these.
data KernelResult
  = -- | Each "worker" in the kernel returns this.
    -- Whether this is a result-per-thread or a
    -- result-per-block depends on where the 'SegOp' occurs.
    Returns ResultManifest Certs SubExp
  | TileReturns
      Certs
      [(SubExp, SubExp)] -- Total/tile for each dimension
      VName -- Tile written by this worker.
      -- The TileReturns must not expect more than one
      -- result to be written per physical thread.
  | RegTileReturns
      Certs
      -- For each dim of result:
      [ ( SubExp, -- size of this dim.
          SubExp, -- block tile size for this dim.
          SubExp -- reg tile size for this dim.
        )
      ]
      VName -- Tile returned by this thread/block.
  deriving (KernelResult -> KernelResult -> Bool
(KernelResult -> KernelResult -> Bool)
-> (KernelResult -> KernelResult -> Bool) -> Eq KernelResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: KernelResult -> KernelResult -> Bool
== :: KernelResult -> KernelResult -> Bool
$c/= :: KernelResult -> KernelResult -> Bool
/= :: KernelResult -> KernelResult -> Bool
Eq, Int -> KernelResult -> ShowS
[KernelResult] -> ShowS
KernelResult -> String
(Int -> KernelResult -> ShowS)
-> (KernelResult -> String)
-> ([KernelResult] -> ShowS)
-> Show KernelResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KernelResult -> ShowS
showsPrec :: Int -> KernelResult -> ShowS
$cshow :: KernelResult -> String
show :: KernelResult -> String
$cshowList :: [KernelResult] -> ShowS
showList :: [KernelResult] -> ShowS
Show, Eq KernelResult
Eq KernelResult =>
(KernelResult -> KernelResult -> Ordering)
-> (KernelResult -> KernelResult -> Bool)
-> (KernelResult -> KernelResult -> Bool)
-> (KernelResult -> KernelResult -> Bool)
-> (KernelResult -> KernelResult -> Bool)
-> (KernelResult -> KernelResult -> KernelResult)
-> (KernelResult -> KernelResult -> KernelResult)
-> Ord KernelResult
KernelResult -> KernelResult -> Bool
KernelResult -> KernelResult -> Ordering
KernelResult -> KernelResult -> KernelResult
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 :: KernelResult -> KernelResult -> Ordering
compare :: KernelResult -> KernelResult -> Ordering
$c< :: KernelResult -> KernelResult -> Bool
< :: KernelResult -> KernelResult -> Bool
$c<= :: KernelResult -> KernelResult -> Bool
<= :: KernelResult -> KernelResult -> Bool
$c> :: KernelResult -> KernelResult -> Bool
> :: KernelResult -> KernelResult -> Bool
$c>= :: KernelResult -> KernelResult -> Bool
>= :: KernelResult -> KernelResult -> Bool
$cmax :: KernelResult -> KernelResult -> KernelResult
max :: KernelResult -> KernelResult -> KernelResult
$cmin :: KernelResult -> KernelResult -> KernelResult
min :: KernelResult -> KernelResult -> KernelResult
Ord)

instance IsResult KernelResult where
  resAliases :: KernelResult -> Names
resAliases = SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn (SubExp -> Names)
-> (KernelResult -> SubExp) -> KernelResult -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. KernelResult -> SubExp
kernelResultSubExp

-- | Get the certs for this 'KernelResult'.
kernelResultCerts :: KernelResult -> Certs
kernelResultCerts :: KernelResult -> Certs
kernelResultCerts (Returns ResultManifest
_ Certs
cs SubExp
_) = Certs
cs
kernelResultCerts (TileReturns Certs
cs [(SubExp, SubExp)]
_ VName
_) = Certs
cs
kernelResultCerts (RegTileReturns Certs
cs [(SubExp, SubExp, SubExp)]
_ VName
_) = Certs
cs

-- | Get the root t'SubExp' corresponding values for a 'KernelResult'.
kernelResultSubExp :: KernelResult -> SubExp
kernelResultSubExp :: KernelResult -> SubExp
kernelResultSubExp (Returns ResultManifest
_ Certs
_ SubExp
se) = SubExp
se
kernelResultSubExp (TileReturns Certs
_ [(SubExp, SubExp)]
_ VName
v) = VName -> SubExp
Var VName
v
kernelResultSubExp (RegTileReturns Certs
_ [(SubExp, SubExp, SubExp)]
_ VName
v) = VName -> SubExp
Var VName
v

instance FreeIn KernelResult where
  freeIn' :: KernelResult -> FV
freeIn' (Returns ResultManifest
_ Certs
cs SubExp
what) = Certs -> FV
forall a. FreeIn a => a -> FV
freeIn' Certs
cs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn' SubExp
what
  freeIn' (TileReturns Certs
cs [(SubExp, SubExp)]
dims VName
v) =
    Certs -> FV
forall a. FreeIn a => a -> FV
freeIn' Certs
cs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [(SubExp, SubExp)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [(SubExp, SubExp)]
dims FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v
  freeIn' (RegTileReturns Certs
cs [(SubExp, SubExp, SubExp)]
dims_n_tiles VName
v) =
    Certs -> FV
forall a. FreeIn a => a -> FV
freeIn' Certs
cs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [(SubExp, SubExp, SubExp)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [(SubExp, SubExp, SubExp)]
dims_n_tiles FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v

instance Substitute KernelResult where
  substituteNames :: Map VName VName -> KernelResult -> KernelResult
substituteNames Map VName VName
subst (Returns ResultManifest
manifest Certs
cs SubExp
se) =
    ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
manifest (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst Certs
cs) (Map VName VName -> SubExp -> SubExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst SubExp
se)
  substituteNames Map VName VName
subst (TileReturns Certs
cs [(SubExp, SubExp)]
dims VName
v) =
    Certs -> [(SubExp, SubExp)] -> VName -> KernelResult
TileReturns
      (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst Certs
cs)
      (Map VName VName -> [(SubExp, SubExp)] -> [(SubExp, SubExp)]
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst [(SubExp, SubExp)]
dims)
      (Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst VName
v)
  substituteNames Map VName VName
subst (RegTileReturns Certs
cs [(SubExp, SubExp, SubExp)]
dims_n_tiles VName
v) =
    Certs -> [(SubExp, SubExp, SubExp)] -> VName -> KernelResult
RegTileReturns
      (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst Certs
cs)
      (Map VName VName
-> [(SubExp, SubExp, SubExp)] -> [(SubExp, SubExp, SubExp)]
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst [(SubExp, SubExp, SubExp)]
dims_n_tiles)
      (Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst VName
v)

instance Rename KernelResult where
  rename :: KernelResult -> RenameM KernelResult
rename = KernelResult -> RenameM KernelResult
forall a. Substitute a => a -> RenameM a
substituteRename

checkKernelResult ::
  (TC.Checkable rep) =>
  KernelResult ->
  Type ->
  TC.TypeM rep ()
checkKernelResult :: forall rep.
Checkable rep =>
KernelResult -> TypeBase Shape NoMode -> TypeM rep ()
checkKernelResult (Returns ResultManifest
_ Certs
cs SubExp
what) TypeBase Shape NoMode
t = do
  Certs -> TypeM rep ()
forall rep. Checkable rep => Certs -> TypeM rep ()
TC.checkCerts Certs
cs
  TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require TypeBase Shape NoMode
t SubExp
what
checkKernelResult (TileReturns Certs
cs [(SubExp, SubExp)]
dims VName
v) TypeBase Shape NoMode
t = do
  Certs -> TypeM rep ()
forall rep. Checkable rep => Certs -> TypeM rep ()
TC.checkCerts Certs
cs
  [(SubExp, SubExp)]
-> ((SubExp, SubExp) -> TypeM rep ()) -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(SubExp, SubExp)]
dims (((SubExp, SubExp) -> TypeM rep ()) -> TypeM rep ())
-> ((SubExp, SubExp) -> TypeM rep ()) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
dim, SubExp
tile) -> do
    TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) SubExp
dim
    TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) SubExp
tile
  vt <- VName -> TypeM rep (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
  unless (vt == t `arrayOfShape` Shape (map snd dims)) $
    TC.bad $
      TC.TypeError $
        "Invalid type for TileReturns " <> prettyText v
checkKernelResult (RegTileReturns Certs
cs [(SubExp, SubExp, SubExp)]
dims_n_tiles VName
arr) TypeBase Shape NoMode
t = do
  Certs -> TypeM rep ()
forall rep. Checkable rep => Certs -> TypeM rep ()
TC.checkCerts Certs
cs
  (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (TypeBase Shape NoMode -> SubExp -> TypeM rep ())
-> TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) [SubExp]
dims
  (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (TypeBase Shape NoMode -> SubExp -> TypeM rep ())
-> TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) [SubExp]
blk_tiles
  (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (TypeBase Shape NoMode -> SubExp -> TypeM rep ())
-> TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) [SubExp]
reg_tiles

  -- assert that arr is of element type t and shape (rev outer_tiles ++ reg_tiles)
  arr_t <- VName -> TypeM rep (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
arr
  unless (arr_t == expected) $
    TC.bad . TC.TypeError $
      "Invalid type for TileReturns. Expected:\n  "
        <> prettyText expected
        <> ",\ngot:\n  "
        <> prettyText arr_t
  where
    ([SubExp]
dims, [SubExp]
blk_tiles, [SubExp]
reg_tiles) = [(SubExp, SubExp, SubExp)] -> ([SubExp], [SubExp], [SubExp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(SubExp, SubExp, SubExp)]
dims_n_tiles
    expected :: TypeBase Shape NoMode
expected = TypeBase Shape NoMode
t TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
`arrayOfShape` [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp]
blk_tiles [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> [SubExp]
reg_tiles)

checkKernelBody ::
  (TC.Checkable rep) =>
  [Type] ->
  KernelBody (Aliases rep) ->
  TC.TypeM rep ()
checkKernelBody :: forall rep.
Checkable rep =>
[TypeBase Shape NoMode] -> KernelBody (Aliases rep) -> TypeM rep ()
checkKernelBody [TypeBase Shape NoMode]
ts (Body (BodyAliasing
_, BodyDec rep
dec) Stms (Aliases rep)
stms [KernelResult]
kres) = do
  BodyDec rep -> TypeM rep ()
forall rep. Checkable rep => BodyDec rep -> TypeM rep ()
TC.checkBodyDec BodyDec rep
dec
  Stms (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Stms (Aliases rep) -> TypeM rep a -> TypeM rep a
TC.checkStms Stms (Aliases rep)
stms (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
    Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoMode]
ts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [KernelResult] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [KernelResult]
kres) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
      ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TC.TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Text
"Kernel return type is "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [TypeBase Shape NoMode] -> Text
forall a. Pretty a => [a] -> Text
prettyTuple [TypeBase Shape NoMode]
ts
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", but body returns "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
forall a. Pretty a => a -> Text
prettyText ([KernelResult] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [KernelResult]
kres)
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" values."
    (KernelResult -> TypeBase Shape NoMode -> TypeM rep ())
-> [KernelResult] -> [TypeBase Shape NoMode] -> TypeM rep ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ KernelResult -> TypeBase Shape NoMode -> TypeM rep ()
forall rep.
Checkable rep =>
KernelResult -> TypeBase Shape NoMode -> TypeM rep ()
checkKernelResult [KernelResult]
kres [TypeBase Shape NoMode]
ts

instance (PrettyRep rep) => Pretty (KernelBody rep) where
  pretty :: forall ann. KernelBody rep -> Doc ann
pretty (Body BodyDec rep
_ Stms rep
stms [KernelResult]
res) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.stack ((Stm rep -> Doc ann) -> [Stm rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Stm rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Stm rep -> Doc ann
pretty (Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms))
      Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"return"
      Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.braces ([Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.commastack ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ (KernelResult -> Doc ann) -> [KernelResult] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map KernelResult -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. KernelResult -> Doc ann
pretty [KernelResult]
res)

certAnnots :: Certs -> [Doc ann]
certAnnots :: forall ann. Certs -> [Doc ann]
certAnnots Certs
cs
  | Certs
cs Certs -> Certs -> Bool
forall a. Eq a => a -> a -> Bool
== Certs
forall a. Monoid a => a
mempty = []
  | Bool
otherwise = [Certs -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Certs -> Doc ann
pretty Certs
cs]

instance Pretty KernelResult where
  pretty :: forall ann. KernelResult -> Doc ann
pretty (Returns ResultManifest
ResultNoSimplify Certs
cs SubExp
what) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
hsep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Certs -> [Doc ann]
forall ann. Certs -> [Doc ann]
certAnnots Certs
cs [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. Semigroup a => a -> a -> a
<> [Doc ann
"returns (manifest)" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
what]
  pretty (Returns ResultManifest
ResultPrivate Certs
cs SubExp
what) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
hsep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Certs -> [Doc ann]
forall ann. Certs -> [Doc ann]
certAnnots Certs
cs [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. Semigroup a => a -> a -> a
<> [Doc ann
"returns (private)" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
what]
  pretty (Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
what) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
hsep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Certs -> [Doc ann]
forall ann. Certs -> [Doc ann]
certAnnots Certs
cs [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. Semigroup a => a -> a -> a
<> [Doc ann
"returns" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
what]
  pretty (TileReturns Certs
cs [(SubExp, SubExp)]
dims VName
v) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
hsep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Certs -> [Doc ann]
forall ann. Certs -> [Doc ann]
certAnnots Certs
cs [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. Semigroup a => a -> a -> a
<> [Doc ann
"tile" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
apply (((SubExp, SubExp) -> Doc ann) -> [(SubExp, SubExp)] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, SubExp) -> Doc ann
forall {a} {a} {ann}. (Pretty a, Pretty a) => (a, a) -> Doc ann
onDim [(SubExp, SubExp)]
dims) Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
v]
    where
      onDim :: (a, a) -> Doc ann
onDim (a
dim, a
tile) = a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
dim Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann
"/" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
tile
  pretty (RegTileReturns Certs
cs [(SubExp, SubExp, SubExp)]
dims_n_tiles VName
v) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
hsep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Certs -> [Doc ann]
forall ann. Certs -> [Doc ann]
certAnnots Certs
cs [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. Semigroup a => a -> a -> a
<> [Doc ann
"blkreg_tile" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
apply (((SubExp, SubExp, SubExp) -> Doc ann)
-> [(SubExp, SubExp, SubExp)] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, SubExp, SubExp) -> Doc ann
forall {a} {a} {a} {ann}.
(Pretty a, Pretty a, Pretty a) =>
(a, a, a) -> Doc ann
onDim [(SubExp, SubExp, SubExp)]
dims_n_tiles) Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
v]
    where
      onDim :: (a, a, a) -> Doc ann
onDim (a
dim, a
blk_tile, a
reg_tile) =
        a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
dim Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann
"/" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
blk_tile Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann
"*" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
reg_tile)

-- | Index space of a 'SegOp'.
data SegSpace = SegSpace
  { -- | Flat physical index corresponding to the
    -- dimensions (at code generation used for a
    -- thread ID or similar).
    SegSpace -> VName
segFlat :: VName,
    SegSpace -> [(VName, SubExp)]
unSegSpace :: [(VName, SubExp)]
  }
  deriving (SegSpace -> SegSpace -> Bool
(SegSpace -> SegSpace -> Bool)
-> (SegSpace -> SegSpace -> Bool) -> Eq SegSpace
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SegSpace -> SegSpace -> Bool
== :: SegSpace -> SegSpace -> Bool
$c/= :: SegSpace -> SegSpace -> Bool
/= :: SegSpace -> SegSpace -> Bool
Eq, Eq SegSpace
Eq SegSpace =>
(SegSpace -> SegSpace -> Ordering)
-> (SegSpace -> SegSpace -> Bool)
-> (SegSpace -> SegSpace -> Bool)
-> (SegSpace -> SegSpace -> Bool)
-> (SegSpace -> SegSpace -> Bool)
-> (SegSpace -> SegSpace -> SegSpace)
-> (SegSpace -> SegSpace -> SegSpace)
-> Ord SegSpace
SegSpace -> SegSpace -> Bool
SegSpace -> SegSpace -> Ordering
SegSpace -> SegSpace -> SegSpace
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 :: SegSpace -> SegSpace -> Ordering
compare :: SegSpace -> SegSpace -> Ordering
$c< :: SegSpace -> SegSpace -> Bool
< :: SegSpace -> SegSpace -> Bool
$c<= :: SegSpace -> SegSpace -> Bool
<= :: SegSpace -> SegSpace -> Bool
$c> :: SegSpace -> SegSpace -> Bool
> :: SegSpace -> SegSpace -> Bool
$c>= :: SegSpace -> SegSpace -> Bool
>= :: SegSpace -> SegSpace -> Bool
$cmax :: SegSpace -> SegSpace -> SegSpace
max :: SegSpace -> SegSpace -> SegSpace
$cmin :: SegSpace -> SegSpace -> SegSpace
min :: SegSpace -> SegSpace -> SegSpace
Ord, Int -> SegSpace -> ShowS
[SegSpace] -> ShowS
SegSpace -> String
(Int -> SegSpace -> ShowS)
-> (SegSpace -> String) -> ([SegSpace] -> ShowS) -> Show SegSpace
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SegSpace -> ShowS
showsPrec :: Int -> SegSpace -> ShowS
$cshow :: SegSpace -> String
show :: SegSpace -> String
$cshowList :: [SegSpace] -> ShowS
showList :: [SegSpace] -> ShowS
Show)

-- | The sizes spanned by the indexes of the 'SegSpace'.
segSpaceDims :: SegSpace -> [SubExp]
segSpaceDims :: SegSpace -> [SubExp]
segSpaceDims (SegSpace VName
_ [(VName, SubExp)]
space) = ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
space

-- | A 'Scope' containing all the identifiers brought into scope by
-- this 'SegSpace'.
scopeOfSegSpace :: SegSpace -> Scope rep
scopeOfSegSpace :: forall rep. SegSpace -> Scope rep
scopeOfSegSpace (SegSpace VName
phys [(VName, SubExp)]
space) =
  [(VName, NameInfo rep)] -> Map VName (NameInfo rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, NameInfo rep)] -> Map VName (NameInfo rep))
-> [(VName, NameInfo rep)] -> Map VName (NameInfo rep)
forall a b. (a -> b) -> a -> b
$ (VName -> (VName, NameInfo rep))
-> [VName] -> [(VName, NameInfo rep)]
forall a b. (a -> b) -> [a] -> [b]
map (,IntType -> NameInfo rep
forall rep. IntType -> NameInfo rep
IndexName IntType
Int64) (VName
phys VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: ((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)]
space)

checkSegSpace :: (TC.Checkable rep) => SegSpace -> TC.TypeM rep ()
checkSegSpace :: forall rep. Checkable rep => SegSpace -> TypeM rep ()
checkSegSpace (SegSpace VName
_ [(VName, SubExp)]
dims) =
  ((VName, SubExp) -> TypeM rep ())
-> [(VName, SubExp)] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) (SubExp -> TypeM rep ())
-> ((VName, SubExp) -> SubExp) -> (VName, SubExp) -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd) [(VName, SubExp)]
dims

-- | A 'SegOp' is semantically a perfectly nested stack of maps, on
-- top of some bottommost computation (scalar computation, reduction,
-- scan, or histogram).  The 'SegSpace' encodes the original map
-- structure.
--
-- All 'SegOp's are parameterised by the representation of their body,
-- as well as a *level*.  The *level* is a representation-specific bit
-- of information.  For example, in GPU backends, it is used to
-- indicate whether the 'SegOp' is expected to run at the thread-level
-- or the block-level.
--
-- The type list is usually the type of the element returned by a
-- single thread. The result of the SegOp is then an array of that
-- type, with the shape of the 'SegSpace' prepended.
data SegOp lvl rep
  = SegMap lvl SegSpace [Type] (KernelBody rep)
  | -- | The KernelSpace must always have at least two dimensions,
    -- implying that the result of a SegRed is always an array.
    SegRed lvl SegSpace [Type] (KernelBody rep) [SegBinOp rep]
  | SegScan lvl SegSpace [Type] (KernelBody rep) [SegBinOp rep] (SegPostOp rep)
  | SegHist lvl SegSpace [Type] (KernelBody rep) [HistOp rep]
  deriving (SegOp lvl rep -> SegOp lvl rep -> Bool
(SegOp lvl rep -> SegOp lvl rep -> Bool)
-> (SegOp lvl rep -> SegOp lvl rep -> Bool) -> Eq (SegOp lvl rep)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall lvl rep.
(RepTypes rep, Eq lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
$c== :: forall lvl rep.
(RepTypes rep, Eq lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
== :: SegOp lvl rep -> SegOp lvl rep -> Bool
$c/= :: forall lvl rep.
(RepTypes rep, Eq lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
/= :: SegOp lvl rep -> SegOp lvl rep -> Bool
Eq, Eq (SegOp lvl rep)
Eq (SegOp lvl rep) =>
(SegOp lvl rep -> SegOp lvl rep -> Ordering)
-> (SegOp lvl rep -> SegOp lvl rep -> Bool)
-> (SegOp lvl rep -> SegOp lvl rep -> Bool)
-> (SegOp lvl rep -> SegOp lvl rep -> Bool)
-> (SegOp lvl rep -> SegOp lvl rep -> Bool)
-> (SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep)
-> (SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep)
-> Ord (SegOp lvl rep)
SegOp lvl rep -> SegOp lvl rep -> Bool
SegOp lvl rep -> SegOp lvl rep -> Ordering
SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep
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
forall lvl rep. (RepTypes rep, Ord lvl) => Eq (SegOp lvl rep)
forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Ordering
forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep
$ccompare :: forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Ordering
compare :: SegOp lvl rep -> SegOp lvl rep -> Ordering
$c< :: forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
< :: SegOp lvl rep -> SegOp lvl rep -> Bool
$c<= :: forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
<= :: SegOp lvl rep -> SegOp lvl rep -> Bool
$c> :: forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
> :: SegOp lvl rep -> SegOp lvl rep -> Bool
$c>= :: forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> Bool
>= :: SegOp lvl rep -> SegOp lvl rep -> Bool
$cmax :: forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep
max :: SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep
$cmin :: forall lvl rep.
(RepTypes rep, Ord lvl) =>
SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep
min :: SegOp lvl rep -> SegOp lvl rep -> SegOp lvl rep
Ord, Int -> SegOp lvl rep -> ShowS
[SegOp lvl rep] -> ShowS
SegOp lvl rep -> String
(Int -> SegOp lvl rep -> ShowS)
-> (SegOp lvl rep -> String)
-> ([SegOp lvl rep] -> ShowS)
-> Show (SegOp lvl rep)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall lvl rep.
(RepTypes rep, Show lvl) =>
Int -> SegOp lvl rep -> ShowS
forall lvl rep.
(RepTypes rep, Show lvl) =>
[SegOp lvl rep] -> ShowS
forall lvl rep. (RepTypes rep, Show lvl) => SegOp lvl rep -> String
$cshowsPrec :: forall lvl rep.
(RepTypes rep, Show lvl) =>
Int -> SegOp lvl rep -> ShowS
showsPrec :: Int -> SegOp lvl rep -> ShowS
$cshow :: forall lvl rep. (RepTypes rep, Show lvl) => SegOp lvl rep -> String
show :: SegOp lvl rep -> String
$cshowList :: forall lvl rep.
(RepTypes rep, Show lvl) =>
[SegOp lvl rep] -> ShowS
showList :: [SegOp lvl rep] -> ShowS
Show)

-- | The level of a 'SegOp'.
segLevel :: SegOp lvl rep -> lvl
segLevel :: forall lvl rep. SegOp lvl rep -> lvl
segLevel (SegMap lvl
lvl SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
_) = lvl
lvl
segLevel (SegRed lvl
lvl SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
_ [SegBinOp rep]
_) = lvl
lvl
segLevel (SegScan lvl
lvl SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
_ [SegBinOp rep]
_ SegPostOp rep
_) = lvl
lvl
segLevel (SegHist lvl
lvl SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
_ [HistOp rep]
_) = lvl
lvl

-- | The space of a 'SegOp'.
segSpace :: SegOp lvl rep -> SegSpace
segSpace :: forall lvl rep. SegOp lvl rep -> SegSpace
segSpace (SegMap lvl
_ SegSpace
lvl [TypeBase Shape NoMode]
_ KernelBody rep
_) = SegSpace
lvl
segSpace (SegRed lvl
_ SegSpace
lvl [TypeBase Shape NoMode]
_ KernelBody rep
_ [SegBinOp rep]
_) = SegSpace
lvl
segSpace (SegScan lvl
_ SegSpace
lvl [TypeBase Shape NoMode]
_ KernelBody rep
_ [SegBinOp rep]
_ SegPostOp rep
_) = SegSpace
lvl
segSpace (SegHist lvl
_ SegSpace
lvl [TypeBase Shape NoMode]
_ KernelBody rep
_ [HistOp rep]
_) = SegSpace
lvl

-- | The body of a 'SegOp'.
segBody :: SegOp lvl rep -> KernelBody rep
segBody :: forall lvl rep. SegOp lvl rep -> KernelBody rep
segBody SegOp lvl rep
segop =
  case SegOp lvl rep
segop of
    SegMap lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body -> KernelBody rep
body
    SegRed lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body [SegBinOp rep]
_ -> KernelBody rep
body
    SegScan lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body [SegBinOp rep]
_ SegPostOp rep
_ -> KernelBody rep
body
    SegHist lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body [HistOp rep]
_ -> KernelBody rep
body

segResultShape :: SegSpace -> Type -> KernelResult -> Type
segResultShape :: SegSpace
-> TypeBase Shape NoMode -> KernelResult -> TypeBase Shape NoMode
segResultShape SegSpace
space TypeBase Shape NoMode
t Returns {} =
  (SubExp -> TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> TypeBase Shape NoMode -> [SubExp] -> TypeBase Shape NoMode
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((TypeBase Shape NoMode -> SubExp -> TypeBase Shape NoMode)
-> SubExp -> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall a b c. (a -> b -> c) -> b -> a -> c
flip TypeBase Shape NoMode -> SubExp -> TypeBase Shape NoMode
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow) TypeBase Shape NoMode
t ([SubExp] -> TypeBase Shape NoMode)
-> [SubExp] -> TypeBase Shape NoMode
forall a b. (a -> b) -> a -> b
$ SegSpace -> [SubExp]
segSpaceDims SegSpace
space
segResultShape SegSpace
_ TypeBase Shape NoMode
t (TileReturns Certs
_ [(SubExp, SubExp)]
dims VName
_) =
  TypeBase Shape NoMode
t TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
`arrayOfShape` [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape (((SubExp, SubExp) -> SubExp) -> [(SubExp, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, SubExp) -> SubExp
forall a b. (a, b) -> a
fst [(SubExp, SubExp)]
dims)
segResultShape SegSpace
_ TypeBase Shape NoMode
t (RegTileReturns Certs
_ [(SubExp, SubExp, SubExp)]
dims_n_tiles VName
_) =
  TypeBase Shape NoMode
t TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
`arrayOfShape` [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape (((SubExp, SubExp, SubExp) -> SubExp)
-> [(SubExp, SubExp, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (\(SubExp
dim, SubExp
_, SubExp
_) -> SubExp
dim) [(SubExp, SubExp, SubExp)]
dims_n_tiles)

-- | The return type of a 'SegOp'.
segOpType :: SegOp lvl rep -> [Type]
segOpType :: forall lvl rep. SegOp lvl rep -> [TypeBase Shape NoMode]
segOpType (SegMap lvl
_ SegSpace
space [TypeBase Shape NoMode]
ts KernelBody rep
kbody) =
  (TypeBase Shape NoMode -> KernelResult -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> [TypeBase Shape NoMode]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (SegSpace
-> TypeBase Shape NoMode -> KernelResult -> TypeBase Shape NoMode
segResultShape SegSpace
space) [TypeBase Shape NoMode]
ts ([KernelResult] -> [TypeBase Shape NoMode])
-> [KernelResult] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$ KernelBody rep -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody rep
kbody
segOpType (SegRed lvl
_ SegSpace
space [TypeBase Shape NoMode]
ts KernelBody rep
kbody [SegBinOp rep]
reds) =
  [TypeBase Shape NoMode]
red_ts
    [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ (TypeBase Shape NoMode -> KernelResult -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> [TypeBase Shape NoMode]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
      (SegSpace
-> TypeBase Shape NoMode -> KernelResult -> TypeBase Shape NoMode
segResultShape SegSpace
space)
      [TypeBase Shape NoMode]
map_ts
      (Int -> [KernelResult] -> [KernelResult]
forall a. Int -> [a] -> [a]
drop ([TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoMode]
red_ts) ([KernelResult] -> [KernelResult])
-> [KernelResult] -> [KernelResult]
forall a b. (a -> b) -> a -> b
$ KernelBody rep -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody rep
kbody)
  where
    map_ts :: [TypeBase Shape NoMode]
map_ts = Int -> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Int -> [a] -> [a]
drop ([TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoMode]
red_ts) [TypeBase Shape NoMode]
ts
    segment_dims :: [SubExp]
segment_dims = [SubExp] -> [SubExp]
forall a. HasCallStack => [a] -> [a]
init ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [SubExp]
segSpaceDims SegSpace
space
    red_ts :: [TypeBase Shape NoMode]
red_ts = do
      op <- [SegBinOp rep]
reds
      let shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
segment_dims Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> SegBinOp rep -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp rep
op
      map (`arrayOfShape` shape) (lambdaReturnType $ segBinOpLambda op)
segOpType (SegScan lvl
_ SegSpace
space [TypeBase Shape NoMode]
_ts KernelBody rep
_kbody [SegBinOp rep]
_scans SegPostOp rep
post_op) =
  (TypeBase Shape NoMode -> [SubExp] -> TypeBase Shape NoMode)
-> [SubExp] -> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SubExp -> TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> TypeBase Shape NoMode -> [SubExp] -> TypeBase Shape NoMode
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((TypeBase Shape NoMode -> SubExp -> TypeBase Shape NoMode)
-> SubExp -> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall a b c. (a -> b -> c) -> b -> a -> c
flip TypeBase Shape NoMode -> SubExp -> TypeBase Shape NoMode
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow)) [SubExp]
segment_dims (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TypeBase Shape NoMode]
lam_ts
  where
    lam_ts :: [TypeBase Shape NoMode]
lam_ts = GLambda rep (TypeBase Shape NoMode) -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (GLambda rep (TypeBase Shape NoMode) -> [TypeBase Shape NoMode])
-> GLambda rep (TypeBase Shape NoMode) -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$ SegPostOp rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp rep
post_op
    segment_dims :: [SubExp]
segment_dims = SegSpace -> [SubExp]
segSpaceDims SegSpace
space
segOpType (SegHist lvl
_ SegSpace
space [TypeBase Shape NoMode]
_ KernelBody rep
_ [HistOp rep]
ops) = do
  op <- [HistOp rep]
ops
  let shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
segment_dims Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> HistOp rep -> Shape
forall rep. HistOp rep -> Shape
histShape HistOp rep
op Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> HistOp rep -> Shape
forall rep. HistOp rep -> Shape
histOpShape HistOp rep
op
  map (`arrayOfShape` shape) (lambdaReturnType $ histOp op)
  where
    dims :: [SubExp]
dims = SegSpace -> [SubExp]
segSpaceDims SegSpace
space
    segment_dims :: [SubExp]
segment_dims = [SubExp] -> [SubExp]
forall a. HasCallStack => [a] -> [a]
init [SubExp]
dims

instance TypedOp (SegOp lvl) where
  opType :: forall rep (m :: * -> *).
HasScope rep m =>
SegOp lvl rep -> m [ExtType]
opType = [ExtType] -> m [ExtType]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExtType] -> m [ExtType])
-> (SegOp lvl rep -> [ExtType]) -> SegOp lvl rep -> m [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [TypeBase Shape NoMode] -> [ExtType]
forall o. [TypeBase Shape o] -> [TypeBase ExtShape o]
staticShapes ([TypeBase Shape NoMode] -> [ExtType])
-> (SegOp lvl rep -> [TypeBase Shape NoMode])
-> SegOp lvl rep
-> [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegOp lvl rep -> [TypeBase Shape NoMode]
forall lvl rep. SegOp lvl rep -> [TypeBase Shape NoMode]
segOpType

instance (ASTConstraints lvl) => AliasedOp (SegOp lvl) where
  opAliases :: forall rep. Aliased rep => SegOp lvl rep -> [Names]
opAliases = (TypeBase Shape NoMode -> Names)
-> [TypeBase Shape NoMode] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (Names -> TypeBase Shape NoMode -> Names
forall a b. a -> b -> a
const Names
forall a. Monoid a => a
mempty) ([TypeBase Shape NoMode] -> [Names])
-> (SegOp lvl rep -> [TypeBase Shape NoMode])
-> SegOp lvl rep
-> [Names]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegOp lvl rep -> [TypeBase Shape NoMode]
forall lvl rep. SegOp lvl rep -> [TypeBase Shape NoMode]
segOpType

  consumedInOp :: forall rep. Aliased rep => SegOp lvl rep -> Names
consumedInOp (SegMap lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
kbody) =
    KernelBody rep -> Names
forall res. GBody rep res -> Names
forall rep res. Aliased rep => GBody rep res -> Names
consumedInBody KernelBody rep
kbody
  consumedInOp (SegRed lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
kbody [SegBinOp rep]
_) =
    KernelBody rep -> Names
forall res. GBody rep res -> Names
forall rep res. Aliased rep => GBody rep res -> Names
consumedInBody KernelBody rep
kbody
  consumedInOp (SegScan lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
kbody [SegBinOp rep]
_ SegPostOp rep
post_op) =
    Names
consumed_lam Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> KernelBody rep -> Names
forall res. GBody rep res -> Names
forall rep res. Aliased rep => GBody rep res -> Names
consumedInBody KernelBody rep
kbody
    where
      consumed_lam :: Names
consumed_lam = Lambda rep -> Names
forall rep. Aliased rep => Lambda rep -> Names
consumedByLambda (Lambda rep -> Names) -> Lambda rep -> Names
forall a b. (a -> b) -> a -> b
$ SegPostOp rep -> Lambda rep
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp rep
post_op
  consumedInOp (SegHist lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
kbody [HistOp rep]
ops) =
    [VName] -> Names
namesFromList ((HistOp rep -> [VName]) -> [HistOp rep] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap HistOp rep -> [VName]
forall rep. HistOp rep -> [VName]
histDest [HistOp rep]
ops) Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> KernelBody rep -> Names
forall res. GBody rep res -> Names
forall rep res. Aliased rep => GBody rep res -> Names
consumedInBody KernelBody rep
kbody

-- | Type check a 'SegOp', given a checker for its level.
typeCheckSegOp ::
  (TC.Checkable rep) =>
  (lvl -> TC.TypeM rep ()) ->
  SegOp lvl (Aliases rep) ->
  TC.TypeM rep ()
typeCheckSegOp :: forall rep lvl.
Checkable rep =>
(lvl -> TypeM rep ()) -> SegOp lvl (Aliases rep) -> TypeM rep ()
typeCheckSegOp lvl -> TypeM rep ()
checkLvl (SegMap lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
kbody) = do
  lvl -> TypeM rep ()
checkLvl lvl
lvl
  SegSpace -> TypeM rep ()
forall rep. Checkable rep => SegSpace -> TypeM rep ()
checkSegSpace SegSpace
space
  Scope (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
TC.binding (SegSpace -> Scope (Aliases rep)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ [(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> KernelBody (Aliases rep)
-> TypeM rep ()
forall rep.
Checkable rep =>
[(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> KernelBody (Aliases rep)
-> TypeM rep ()
checkScanRed [] [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
kbody
typeCheckSegOp lvl -> TypeM rep ()
checkLvl (SegRed lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
body [SegBinOp (Aliases rep)]
reds) = do
  lvl -> TypeM rep ()
checkLvl lvl
lvl
  SegSpace -> TypeM rep ()
forall rep. Checkable rep => SegSpace -> TypeM rep ()
checkSegSpace SegSpace
space
  Scope (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
TC.binding (SegSpace -> Scope (Aliases rep)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ [(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> KernelBody (Aliases rep)
-> TypeM rep ()
forall rep.
Checkable rep =>
[(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> KernelBody (Aliases rep)
-> TypeM rep ()
checkScanRed [(Lambda (Aliases rep), [SubExp], Shape)]
reds' [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
body
  where
    reds' :: [(Lambda (Aliases rep), [SubExp], Shape)]
reds' =
      [Lambda (Aliases rep)]
-> [[SubExp]]
-> [Shape]
-> [(Lambda (Aliases rep), [SubExp], Shape)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3
        ((SegBinOp (Aliases rep) -> Lambda (Aliases rep))
-> [SegBinOp (Aliases rep)] -> [Lambda (Aliases rep)]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp (Aliases rep) -> Lambda (Aliases rep)
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda [SegBinOp (Aliases rep)]
reds)
        ((SegBinOp (Aliases rep) -> [SubExp])
-> [SegBinOp (Aliases rep)] -> [[SubExp]]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp (Aliases rep) -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral [SegBinOp (Aliases rep)]
reds)
        ((SegBinOp (Aliases rep) -> Shape)
-> [SegBinOp (Aliases rep)] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp (Aliases rep) -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape [SegBinOp (Aliases rep)]
reds)
typeCheckSegOp lvl -> TypeM rep ()
checkLvl (SegScan lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
body [SegBinOp (Aliases rep)]
scans SegPostOp (Aliases rep)
post_op) = do
  lvl -> TypeM rep ()
checkLvl lvl
lvl
  SegSpace -> TypeM rep ()
forall rep. Checkable rep => SegSpace -> TypeM rep ()
checkSegSpace SegSpace
space
  Scope (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
TC.binding (SegSpace -> Scope (Aliases rep)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
    [(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> KernelBody (Aliases rep)
-> TypeM rep ()
forall rep.
Checkable rep =>
[(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> KernelBody (Aliases rep)
-> TypeM rep ()
checkScanRed [(Lambda (Aliases rep), [SubExp], Shape)]
scans' [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
body
    SegPostOp (Aliases rep)
-> [(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> TypeM rep ()
forall rep.
Checkable rep =>
SegPostOp (Aliases rep)
-> [(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> TypeM rep ()
checkSegPostOp SegPostOp (Aliases rep)
post_op [(Lambda (Aliases rep), [SubExp], Shape)]
scans' [TypeBase Shape NoMode]
ts
  where
    scans' :: [(Lambda (Aliases rep), [SubExp], Shape)]
scans' =
      [Lambda (Aliases rep)]
-> [[SubExp]]
-> [Shape]
-> [(Lambda (Aliases rep), [SubExp], Shape)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3
        ((SegBinOp (Aliases rep) -> Lambda (Aliases rep))
-> [SegBinOp (Aliases rep)] -> [Lambda (Aliases rep)]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp (Aliases rep) -> Lambda (Aliases rep)
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda [SegBinOp (Aliases rep)]
scans)
        ((SegBinOp (Aliases rep) -> [SubExp])
-> [SegBinOp (Aliases rep)] -> [[SubExp]]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp (Aliases rep) -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral [SegBinOp (Aliases rep)]
scans)
        ((SegBinOp (Aliases rep) -> Shape)
-> [SegBinOp (Aliases rep)] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp (Aliases rep) -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape [SegBinOp (Aliases rep)]
scans)
typeCheckSegOp lvl -> TypeM rep ()
checkLvl (SegHist lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
kbody [HistOp (Aliases rep)]
ops) = do
  lvl -> TypeM rep ()
checkLvl lvl
lvl
  SegSpace -> TypeM rep ()
forall rep. Checkable rep => SegSpace -> TypeM rep ()
checkSegSpace SegSpace
space
  (TypeBase Shape NoMode -> TypeM rep ())
-> [TypeBase Shape NoMode] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ TypeBase Shape NoMode -> TypeM rep ()
forall rep o. Checkable rep => TypeBase Shape o -> TypeM rep ()
TC.checkType [TypeBase Shape NoMode]
ts

  Scope (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
TC.binding (SegSpace -> Scope (Aliases rep)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
    nes_ts <- [HistOp (Aliases rep)]
-> (HistOp (Aliases rep) -> TypeM rep [TypeBase Shape NoMode])
-> TypeM rep [[TypeBase Shape NoMode]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [HistOp (Aliases rep)]
ops ((HistOp (Aliases rep) -> TypeM rep [TypeBase Shape NoMode])
 -> TypeM rep [[TypeBase Shape NoMode]])
-> (HistOp (Aliases rep) -> TypeM rep [TypeBase Shape NoMode])
-> TypeM rep [[TypeBase Shape NoMode]]
forall a b. (a -> b) -> a -> b
$ \(HistOp Shape
dest_shape SubExp
rf [VName]
dests [SubExp]
nes Shape
shape Lambda (Aliases rep)
op) -> do
      (SubExp -> TypeM rep ()) -> Shape -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64)) Shape
dest_shape
      TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) SubExp
rf
      nes' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Arg
forall rep. Checkable rep => SubExp -> TypeM rep Arg
TC.checkArg [SubExp]
nes
      mapM_ (TC.require (Prim int64)) $ shapeDims shape

      -- Operator type must match the type of neutral elements.
      let stripVecDims = Int -> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray (Int -> TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> Int -> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall a b. (a -> b) -> a -> b
$ Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
shape
      TC.checkLambda op $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes'
      let nes_t = (Arg -> TypeBase Shape NoMode) -> [Arg] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> TypeBase Shape NoMode
TC.argType [Arg]
nes'
      unless (nes_t == lambdaReturnType op) $
        TC.bad $
          TC.TypeError $
            "SegHist operator has return type "
              <> prettyTuple (lambdaReturnType op)
              <> " but neutral element has type "
              <> prettyTuple nes_t

      -- Arrays must have proper type.
      let dest_shape' = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
segment_dims Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
dest_shape Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
shape
      forM_ (zip nes_t dests) $ \(TypeBase Shape NoMode
t, VName
dest) -> do
        TypeBase Shape NoMode -> VName -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> VName -> TypeM rep ()
TC.requireI (TypeBase Shape NoMode
t TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
`arrayOfShape` Shape
dest_shape') VName
dest
        Names -> TypeM rep ()
forall rep. Checkable rep => Names -> TypeM rep ()
TC.consume (Names -> TypeM rep ()) -> TypeM rep Names -> TypeM rep ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> TypeM rep Names
forall rep. Checkable rep => VName -> TypeM rep Names
TC.lookupAliases VName
dest

      pure $ map (`arrayOfShape` shape) nes_t

    checkKernelBody ts kbody

    -- Return type of bucket function must be an index for each
    -- operation followed by the values to write.
    let bucket_ret_t =
          (HistOp (Aliases rep) -> [TypeBase Shape NoMode])
-> [HistOp (Aliases rep)] -> [TypeBase Shape NoMode]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Int -> TypeBase Shape NoMode -> [TypeBase Shape NoMode]
forall a. Int -> a -> [a]
`replicate` PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) (Int -> [TypeBase Shape NoMode])
-> (HistOp (Aliases rep) -> Int)
-> HistOp (Aliases rep)
-> [TypeBase Shape NoMode]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int)
-> (HistOp (Aliases rep) -> Shape) -> HistOp (Aliases rep) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. HistOp (Aliases rep) -> Shape
forall rep. HistOp rep -> Shape
histShape) [HistOp (Aliases rep)]
ops
            [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ [[TypeBase Shape NoMode]] -> [TypeBase Shape NoMode]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase Shape NoMode]]
nes_ts
    unless (bucket_ret_t == ts) $
      TC.bad $
        TC.TypeError $
          "SegHist body has return type "
            <> prettyTuple ts
            <> " but should have type "
            <> prettyTuple bucket_ret_t
  where
    segment_dims :: [SubExp]
segment_dims = [SubExp] -> [SubExp]
forall a. HasCallStack => [a] -> [a]
init ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [SubExp]
segSpaceDims SegSpace
space

checkSegPostOp ::
  (TC.Checkable rep) =>
  SegPostOp (Aliases rep) ->
  [(Lambda (Aliases rep), [SubExp], Shape)] ->
  [Type] ->
  TC.TypeM rep ()
checkSegPostOp :: forall rep.
Checkable rep =>
SegPostOp (Aliases rep)
-> [(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> TypeM rep ()
checkSegPostOp (SegPostOp Lambda (Aliases rep)
lam) [(Lambda (Aliases rep), [SubExp], Shape)]
ops [TypeBase Shape NoMode]
ts = do
  let ([Shape]
shps, [SubExp]
nes) = [(Shape, SubExp)] -> ([Shape], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Shape, SubExp)] -> ([Shape], [SubExp]))
-> [(Shape, SubExp)] -> ([Shape], [SubExp])
forall a b. (a -> b) -> a -> b
$ ((Lambda (Aliases rep), [SubExp], Shape) -> [(Shape, SubExp)])
-> [(Lambda (Aliases rep), [SubExp], Shape)] -> [(Shape, SubExp)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Lambda (Aliases rep)
_, [SubExp]
a, Shape
shp) -> (Shape
shp,) (SubExp -> (Shape, SubExp)) -> [SubExp] -> [(Shape, SubExp)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExp]
a) [(Lambda (Aliases rep), [SubExp], Shape)]
ops
  nes' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Arg
forall rep. Checkable rep => SubExp -> TypeM rep Arg
TC.checkArg [SubExp]
nes
  let kbody' = (,Names
forall a. Monoid a => a
mempty) (TypeBase Shape NoMode -> Arg) -> [TypeBase Shape NoMode] -> [Arg]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Int -> [a] -> [a]
drop ([Arg] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Arg]
nes') [TypeBase Shape NoMode]
ts

  TC.checkLambda lam $
    map TC.noArgAliases $
      zipWith (\Shape
shp (TypeBase Shape NoMode
t, Names
e) -> (TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
arrayOfShape TypeBase Shape NoMode
t Shape
shp, Names
e)) shps nes' <> kbody'

checkScanRed ::
  (TC.Checkable rep) =>
  [(Lambda (Aliases rep), [SubExp], Shape)] ->
  [Type] ->
  KernelBody (Aliases rep) ->
  TC.TypeM rep ()
checkScanRed :: forall rep.
Checkable rep =>
[(Lambda (Aliases rep), [SubExp], Shape)]
-> [TypeBase Shape NoMode]
-> KernelBody (Aliases rep)
-> TypeM rep ()
checkScanRed [(Lambda (Aliases rep), [SubExp], Shape)]
ops [TypeBase Shape NoMode]
ts KernelBody (Aliases rep)
kbody = do
  (TypeBase Shape NoMode -> TypeM rep ())
-> [TypeBase Shape NoMode] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ TypeBase Shape NoMode -> TypeM rep ()
forall rep o. Checkable rep => TypeBase Shape o -> TypeM rep ()
TC.checkType [TypeBase Shape NoMode]
ts

  ne_ts <- [(Lambda (Aliases rep), [SubExp], Shape)]
-> ((Lambda (Aliases rep), [SubExp], Shape)
    -> TypeM rep [TypeBase Shape NoMode])
-> TypeM rep [[TypeBase Shape NoMode]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Lambda (Aliases rep), [SubExp], Shape)]
ops (((Lambda (Aliases rep), [SubExp], Shape)
  -> TypeM rep [TypeBase Shape NoMode])
 -> TypeM rep [[TypeBase Shape NoMode]])
-> ((Lambda (Aliases rep), [SubExp], Shape)
    -> TypeM rep [TypeBase Shape NoMode])
-> TypeM rep [[TypeBase Shape NoMode]]
forall a b. (a -> b) -> a -> b
$ \(Lambda (Aliases rep)
lam, [SubExp]
nes, Shape
shape) -> do
    (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase Shape NoMode -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
TypeBase Shape NoMode -> SubExp -> TypeM rep ()
TC.require (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64)) ([SubExp] -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
    nes' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Arg
forall rep. Checkable rep => SubExp -> TypeM rep Arg
TC.checkArg [SubExp]
nes

    -- Operator type must match the type of neutral elements.
    TC.checkLambda lam $ map TC.noArgAliases $ nes' ++ nes'
    let nes_t = (Arg -> TypeBase Shape NoMode) -> [Arg] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> TypeBase Shape NoMode
TC.argType [Arg]
nes'

    unless (lambdaReturnType lam == nes_t) $
      TC.bad $
        TC.TypeError "wrong type for operator or neutral elements."

    pure $ map (`arrayOfShape` shape) nes_t

  let expecting = [[TypeBase Shape NoMode]] -> [TypeBase Shape NoMode]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase Shape NoMode]]
ne_ts
      got = Int -> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Int -> [a] -> [a]
take ([TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoMode]
expecting) [TypeBase Shape NoMode]
ts
  unless (expecting == got) $
    TC.bad $
      TC.TypeError $
        "Wrong return for body (does not match neutral elements; expected "
          <> prettyText expecting
          <> "; found "
          <> prettyText got
          <> ")"

  checkKernelBody ts kbody

-- | Like 'Mapper', but just for 'SegOp's.
data SegOpMapper lvl frep trep m = SegOpMapper
  { forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp :: SubExp -> m SubExp,
    forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSegBinOpLambda :: Lambda frep -> m (Lambda trep),
    forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSegPostOpLambda :: Lambda frep -> m (Lambda trep),
    forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
mapOnSegOpBody :: KernelBody frep -> m (KernelBody trep),
    forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> VName -> m VName
mapOnSegOpVName :: VName -> m VName,
    forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> lvl -> m lvl
mapOnSegOpLevel :: lvl -> m lvl
  }

-- | A mapper that simply returns the 'SegOp' verbatim.
identitySegOpMapper :: (Monad m) => SegOpMapper lvl rep rep m
identitySegOpMapper :: forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
identitySegOpMapper =
  SegOpMapper
    { mapOnSegOpSubExp :: SubExp -> m SubExp
mapOnSegOpSubExp = SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSegBinOpLambda :: Lambda rep -> m (Lambda rep)
mapOnSegBinOpLambda = Lambda rep -> m (Lambda rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSegPostOpLambda :: Lambda rep -> m (Lambda rep)
mapOnSegPostOpLambda = Lambda rep -> m (Lambda rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSegOpBody :: KernelBody rep -> m (KernelBody rep)
mapOnSegOpBody = KernelBody rep -> m (KernelBody rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSegOpVName :: VName -> m VName
mapOnSegOpVName = VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSegOpLevel :: lvl -> m lvl
mapOnSegOpLevel = lvl -> m lvl
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    }

mapOnSegSpace ::
  (Monad f) => SegOpMapper lvl frep trep f -> SegSpace -> f SegSpace
mapOnSegSpace :: forall (f :: * -> *) lvl frep trep.
Monad f =>
SegOpMapper lvl frep trep f -> SegSpace -> f SegSpace
mapOnSegSpace SegOpMapper lvl frep trep f
tv (SegSpace VName
phys [(VName, SubExp)]
dims) =
  VName -> [(VName, SubExp)] -> SegSpace
SegSpace
    (VName -> [(VName, SubExp)] -> SegSpace)
-> f VName -> f ([(VName, SubExp)] -> SegSpace)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep f -> VName -> f VName
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> VName -> m VName
mapOnSegOpVName SegOpMapper lvl frep trep f
tv VName
phys
    f ([(VName, SubExp)] -> SegSpace)
-> f [(VName, SubExp)] -> f SegSpace
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((VName, SubExp) -> f (VName, SubExp))
-> [(VName, SubExp)] -> f [(VName, SubExp)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((VName -> f VName)
-> (SubExp -> f SubExp) -> (VName, SubExp) -> f (VName, SubExp)
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> (a, b) -> f (c, d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (SegOpMapper lvl frep trep f -> VName -> f VName
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> VName -> m VName
mapOnSegOpVName SegOpMapper lvl frep trep f
tv) (SegOpMapper lvl frep trep f -> SubExp -> f SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep f
tv)) [(VName, SubExp)]
dims

mapSegBinOp ::
  (Monad m) =>
  SegOpMapper lvl frep trep m ->
  SegBinOp frep ->
  m (SegBinOp trep)
mapSegBinOp :: forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegBinOp frep -> m (SegBinOp trep)
mapSegBinOp SegOpMapper lvl frep trep m
tv (SegBinOp Commutativity
comm Lambda frep
red_op [SubExp]
nes Shape
shape) =
  Commutativity -> Lambda trep -> [SubExp] -> Shape -> SegBinOp trep
forall rep.
Commutativity -> Lambda rep -> [SubExp] -> Shape -> SegBinOp rep
SegBinOp Commutativity
comm
    (Lambda trep -> [SubExp] -> Shape -> SegBinOp trep)
-> m (Lambda trep) -> m ([SubExp] -> Shape -> SegBinOp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSegBinOpLambda SegOpMapper lvl frep trep m
tv Lambda frep
red_op
    m ([SubExp] -> Shape -> SegBinOp trep)
-> m [SubExp] -> m (Shape -> SegBinOp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) [SubExp]
nes
    m (Shape -> SegBinOp trep) -> m Shape -> m (SegBinOp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> m [SubExp] -> m Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape))

mapSegPostOp ::
  (Monad m) =>
  SegOpMapper lvl frep trep m ->
  SegPostOp frep ->
  m (SegPostOp trep)
mapSegPostOp :: forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegPostOp frep -> m (SegPostOp trep)
mapSegPostOp SegOpMapper lvl frep trep m
tv (SegPostOp Lambda frep
lam) =
  Lambda trep -> SegPostOp trep
forall rep. Lambda rep -> SegPostOp rep
SegPostOp (Lambda trep -> SegPostOp trep)
-> m (Lambda trep) -> m (SegPostOp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSegPostOpLambda SegOpMapper lvl frep trep m
tv Lambda frep
lam

-- | Apply a 'SegOpMapper' to the given 'SegOp'.
mapSegOpM ::
  (Monad m) =>
  SegOpMapper lvl frep trep m ->
  SegOp lvl frep ->
  m (SegOp lvl trep)
mapSegOpM :: forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper lvl frep trep m
tv (SegMap lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody frep
body) =
  lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody trep
-> SegOp lvl trep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> SegOp lvl rep
SegMap
    (lvl
 -> SegSpace
 -> [TypeBase Shape NoMode]
 -> KernelBody trep
 -> SegOp lvl trep)
-> m lvl
-> m (SegSpace
      -> [TypeBase Shape NoMode] -> KernelBody trep -> SegOp lvl trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep m -> lvl -> m lvl
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> lvl -> m lvl
mapOnSegOpLevel SegOpMapper lvl frep trep m
tv lvl
lvl
    m (SegSpace
   -> [TypeBase Shape NoMode] -> KernelBody trep -> SegOp lvl trep)
-> m SegSpace
-> m ([TypeBase Shape NoMode] -> KernelBody trep -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m -> SegSpace -> m SegSpace
forall (f :: * -> *) lvl frep trep.
Monad f =>
SegOpMapper lvl frep trep f -> SegSpace -> f SegSpace
mapOnSegSpace SegOpMapper lvl frep trep m
tv SegSpace
space
    m ([TypeBase Shape NoMode] -> KernelBody trep -> SegOp lvl trep)
-> m [TypeBase Shape NoMode]
-> m (KernelBody trep -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> [TypeBase Shape NoMode] -> m [TypeBase Shape NoMode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m
-> TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m
-> TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
mapOnSegOpType SegOpMapper lvl frep trep m
tv) [TypeBase Shape NoMode]
ts
    m (KernelBody trep -> SegOp lvl trep)
-> m (KernelBody trep) -> m (SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
mapOnSegOpBody SegOpMapper lvl frep trep m
tv KernelBody frep
body
mapSegOpM SegOpMapper lvl frep trep m
tv (SegRed lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody frep
body [SegBinOp frep]
reds) =
  lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody trep
-> [SegBinOp trep]
-> SegOp lvl trep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
SegRed
    (lvl
 -> SegSpace
 -> [TypeBase Shape NoMode]
 -> KernelBody trep
 -> [SegBinOp trep]
 -> SegOp lvl trep)
-> m lvl
-> m (SegSpace
      -> [TypeBase Shape NoMode]
      -> KernelBody trep
      -> [SegBinOp trep]
      -> SegOp lvl trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep m -> lvl -> m lvl
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> lvl -> m lvl
mapOnSegOpLevel SegOpMapper lvl frep trep m
tv lvl
lvl
    m (SegSpace
   -> [TypeBase Shape NoMode]
   -> KernelBody trep
   -> [SegBinOp trep]
   -> SegOp lvl trep)
-> m SegSpace
-> m ([TypeBase Shape NoMode]
      -> KernelBody trep -> [SegBinOp trep] -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m -> SegSpace -> m SegSpace
forall (f :: * -> *) lvl frep trep.
Monad f =>
SegOpMapper lvl frep trep f -> SegSpace -> f SegSpace
mapOnSegSpace SegOpMapper lvl frep trep m
tv SegSpace
space
    m ([TypeBase Shape NoMode]
   -> KernelBody trep -> [SegBinOp trep] -> SegOp lvl trep)
-> m [TypeBase Shape NoMode]
-> m (KernelBody trep -> [SegBinOp trep] -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> [TypeBase Shape NoMode] -> m [TypeBase Shape NoMode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SubExp -> m SubExp)
-> TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
forall (m :: * -> *) o.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape o -> m (TypeBase Shape o)
mapOnType ((SubExp -> m SubExp)
 -> TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> (SubExp -> m SubExp)
-> TypeBase Shape NoMode
-> m (TypeBase Shape NoMode)
forall a b. (a -> b) -> a -> b
$ SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) [TypeBase Shape NoMode]
ts
    m (KernelBody trep -> [SegBinOp trep] -> SegOp lvl trep)
-> m (KernelBody trep) -> m ([SegBinOp trep] -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
mapOnSegOpBody SegOpMapper lvl frep trep m
tv KernelBody frep
body
    m ([SegBinOp trep] -> SegOp lvl trep)
-> m [SegBinOp trep] -> m (SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SegBinOp frep -> m (SegBinOp trep))
-> [SegBinOp frep] -> m [SegBinOp trep]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m -> SegBinOp frep -> m (SegBinOp trep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegBinOp frep -> m (SegBinOp trep)
mapSegBinOp SegOpMapper lvl frep trep m
tv) [SegBinOp frep]
reds
mapSegOpM SegOpMapper lvl frep trep m
tv (SegScan lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody frep
body [SegBinOp frep]
scans SegPostOp frep
post_op) =
  lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody trep
-> [SegBinOp trep]
-> SegPostOp trep
-> SegOp lvl trep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegPostOp rep
-> SegOp lvl rep
SegScan
    (lvl
 -> SegSpace
 -> [TypeBase Shape NoMode]
 -> KernelBody trep
 -> [SegBinOp trep]
 -> SegPostOp trep
 -> SegOp lvl trep)
-> m lvl
-> m (SegSpace
      -> [TypeBase Shape NoMode]
      -> KernelBody trep
      -> [SegBinOp trep]
      -> SegPostOp trep
      -> SegOp lvl trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep m -> lvl -> m lvl
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> lvl -> m lvl
mapOnSegOpLevel SegOpMapper lvl frep trep m
tv lvl
lvl
    m (SegSpace
   -> [TypeBase Shape NoMode]
   -> KernelBody trep
   -> [SegBinOp trep]
   -> SegPostOp trep
   -> SegOp lvl trep)
-> m SegSpace
-> m ([TypeBase Shape NoMode]
      -> KernelBody trep
      -> [SegBinOp trep]
      -> SegPostOp trep
      -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m -> SegSpace -> m SegSpace
forall (f :: * -> *) lvl frep trep.
Monad f =>
SegOpMapper lvl frep trep f -> SegSpace -> f SegSpace
mapOnSegSpace SegOpMapper lvl frep trep m
tv SegSpace
space
    m ([TypeBase Shape NoMode]
   -> KernelBody trep
   -> [SegBinOp trep]
   -> SegPostOp trep
   -> SegOp lvl trep)
-> m [TypeBase Shape NoMode]
-> m (KernelBody trep
      -> [SegBinOp trep] -> SegPostOp trep -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> [TypeBase Shape NoMode] -> m [TypeBase Shape NoMode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SubExp -> m SubExp)
-> TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
forall (m :: * -> *) o.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape o -> m (TypeBase Shape o)
mapOnType ((SubExp -> m SubExp)
 -> TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> (SubExp -> m SubExp)
-> TypeBase Shape NoMode
-> m (TypeBase Shape NoMode)
forall a b. (a -> b) -> a -> b
$ SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) [TypeBase Shape NoMode]
ts
    m (KernelBody trep
   -> [SegBinOp trep] -> SegPostOp trep -> SegOp lvl trep)
-> m (KernelBody trep)
-> m ([SegBinOp trep] -> SegPostOp trep -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
mapOnSegOpBody SegOpMapper lvl frep trep m
tv KernelBody frep
body
    m ([SegBinOp trep] -> SegPostOp trep -> SegOp lvl trep)
-> m [SegBinOp trep] -> m (SegPostOp trep -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SegBinOp frep -> m (SegBinOp trep))
-> [SegBinOp frep] -> m [SegBinOp trep]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m -> SegBinOp frep -> m (SegBinOp trep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegBinOp frep -> m (SegBinOp trep)
mapSegBinOp SegOpMapper lvl frep trep m
tv) [SegBinOp frep]
scans
    m (SegPostOp trep -> SegOp lvl trep)
-> m (SegPostOp trep) -> m (SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m -> SegPostOp frep -> m (SegPostOp trep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegPostOp frep -> m (SegPostOp trep)
mapSegPostOp SegOpMapper lvl frep trep m
tv SegPostOp frep
post_op
mapSegOpM SegOpMapper lvl frep trep m
tv (SegHist lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody frep
body [HistOp frep]
ops) =
  lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody trep
-> [HistOp trep]
-> SegOp lvl trep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [HistOp rep]
-> SegOp lvl rep
SegHist
    (lvl
 -> SegSpace
 -> [TypeBase Shape NoMode]
 -> KernelBody trep
 -> [HistOp trep]
 -> SegOp lvl trep)
-> m lvl
-> m (SegSpace
      -> [TypeBase Shape NoMode]
      -> KernelBody trep
      -> [HistOp trep]
      -> SegOp lvl trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep m -> lvl -> m lvl
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> lvl -> m lvl
mapOnSegOpLevel SegOpMapper lvl frep trep m
tv lvl
lvl
    m (SegSpace
   -> [TypeBase Shape NoMode]
   -> KernelBody trep
   -> [HistOp trep]
   -> SegOp lvl trep)
-> m SegSpace
-> m ([TypeBase Shape NoMode]
      -> KernelBody trep -> [HistOp trep] -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m -> SegSpace -> m SegSpace
forall (f :: * -> *) lvl frep trep.
Monad f =>
SegOpMapper lvl frep trep f -> SegSpace -> f SegSpace
mapOnSegSpace SegOpMapper lvl frep trep m
tv SegSpace
space
    m ([TypeBase Shape NoMode]
   -> KernelBody trep -> [HistOp trep] -> SegOp lvl trep)
-> m [TypeBase Shape NoMode]
-> m (KernelBody trep -> [HistOp trep] -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> [TypeBase Shape NoMode] -> m [TypeBase Shape NoMode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SubExp -> m SubExp)
-> TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
forall (m :: * -> *) o.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape o -> m (TypeBase Shape o)
mapOnType ((SubExp -> m SubExp)
 -> TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> (SubExp -> m SubExp)
-> TypeBase Shape NoMode
-> m (TypeBase Shape NoMode)
forall a b. (a -> b) -> a -> b
$ SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) [TypeBase Shape NoMode]
ts
    m (KernelBody trep -> [HistOp trep] -> SegOp lvl trep)
-> m (KernelBody trep) -> m ([HistOp trep] -> SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m
-> KernelBody frep -> m (KernelBody trep)
mapOnSegOpBody SegOpMapper lvl frep trep m
tv KernelBody frep
body
    m ([HistOp trep] -> SegOp lvl trep)
-> m [HistOp trep] -> m (SegOp lvl trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (HistOp frep -> m (HistOp trep))
-> [HistOp frep] -> m [HistOp trep]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM HistOp frep -> m (HistOp trep)
onHistOp [HistOp frep]
ops
  where
    onHistOp :: HistOp frep -> m (HistOp trep)
onHistOp (HistOp Shape
w SubExp
rf [VName]
arrs [SubExp]
nes Shape
shape Lambda frep
op) =
      Shape
-> SubExp
-> [VName]
-> [SubExp]
-> Shape
-> Lambda trep
-> HistOp trep
forall rep.
Shape
-> SubExp
-> [VName]
-> [SubExp]
-> Shape
-> Lambda rep
-> HistOp rep
HistOp
        (Shape
 -> SubExp
 -> [VName]
 -> [SubExp]
 -> Shape
 -> Lambda trep
 -> HistOp trep)
-> m Shape
-> m (SubExp
      -> [VName] -> [SubExp] -> Shape -> Lambda trep -> HistOp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> Shape -> m Shape
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> ShapeBase a -> m (ShapeBase b)
mapM (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) Shape
w
        m (SubExp
   -> [VName] -> [SubExp] -> Shape -> Lambda trep -> HistOp trep)
-> m SubExp
-> m ([VName] -> [SubExp] -> Shape -> Lambda trep -> HistOp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv SubExp
rf
        m ([VName] -> [SubExp] -> Shape -> Lambda trep -> HistOp trep)
-> m [VName] -> m ([SubExp] -> Shape -> Lambda trep -> HistOp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (VName -> m VName) -> [VName] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m -> VName -> m VName
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> VName -> m VName
mapOnSegOpVName SegOpMapper lvl frep trep m
tv) [VName]
arrs
        m ([SubExp] -> Shape -> Lambda trep -> HistOp trep)
-> m [SubExp] -> m (Shape -> Lambda trep -> HistOp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) [SubExp]
nes
        m (Shape -> Lambda trep -> HistOp trep)
-> m Shape -> m (Lambda trep -> HistOp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> m [SubExp] -> m Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape))
        m (Lambda trep -> HistOp trep)
-> m (Lambda trep) -> m (HistOp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSegBinOpLambda SegOpMapper lvl frep trep m
tv Lambda frep
op

mapOnSegOpType ::
  (Monad m) =>
  SegOpMapper lvl frep trep m ->
  Type ->
  m Type
mapOnSegOpType :: forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m
-> TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
mapOnSegOpType SegOpMapper lvl frep trep m
_tv t :: TypeBase Shape NoMode
t@Prim {} = TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase Shape NoMode
t
mapOnSegOpType SegOpMapper lvl frep trep m
tv (Acc VName
acc Shape
ispace [TypeBase Shape NoMode]
ts) =
  VName -> Shape -> [TypeBase Shape NoMode] -> TypeBase Shape NoMode
forall shape o.
VName -> Shape -> [TypeBase Shape NoMode] -> TypeBase shape o
Acc
    (VName
 -> Shape -> [TypeBase Shape NoMode] -> TypeBase Shape NoMode)
-> m VName
-> m (Shape -> [TypeBase Shape NoMode] -> TypeBase Shape NoMode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOpMapper lvl frep trep m -> VName -> m VName
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> VName -> m VName
mapOnSegOpVName SegOpMapper lvl frep trep m
tv VName
acc
    m (Shape -> [TypeBase Shape NoMode] -> TypeBase Shape NoMode)
-> m Shape -> m ([TypeBase Shape NoMode] -> TypeBase Shape NoMode)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> Shape -> m Shape
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> ShapeBase a -> f (ShapeBase b)
traverse (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) Shape
ispace
    m ([TypeBase Shape NoMode] -> TypeBase Shape NoMode)
-> m [TypeBase Shape NoMode] -> m (TypeBase Shape NoMode)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> [TypeBase Shape NoMode] -> m [TypeBase Shape NoMode]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((Shape -> m Shape)
-> (NoMode -> m NoMode)
-> TypeBase Shape NoMode
-> m (TypeBase Shape NoMode)
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse ((SubExp -> m SubExp) -> Shape -> m Shape
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> ShapeBase a -> f (ShapeBase b)
traverse (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv)) NoMode -> m NoMode
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) [TypeBase Shape NoMode]
ts
mapOnSegOpType SegOpMapper lvl frep trep m
tv (Array PrimType
et Shape
shape NoMode
o) =
  PrimType -> Shape -> NoMode -> TypeBase Shape NoMode
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
et (Shape -> NoMode -> TypeBase Shape NoMode)
-> m Shape -> m (NoMode -> TypeBase Shape NoMode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> Shape -> m Shape
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> ShapeBase a -> f (ShapeBase b)
traverse (SegOpMapper lvl frep trep m -> SubExp -> m SubExp
forall lvl frep trep (m :: * -> *).
SegOpMapper lvl frep trep m -> SubExp -> m SubExp
mapOnSegOpSubExp SegOpMapper lvl frep trep m
tv) Shape
shape m (NoMode -> TypeBase Shape NoMode)
-> m NoMode -> m (TypeBase Shape NoMode)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> NoMode -> m NoMode
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NoMode
o
mapOnSegOpType SegOpMapper lvl frep trep m
_tv (Mem Space
s) = TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Shape NoMode -> m (TypeBase Shape NoMode))
-> TypeBase Shape NoMode -> m (TypeBase Shape NoMode)
forall a b. (a -> b) -> a -> b
$ Space -> TypeBase Shape NoMode
forall shape o. Space -> TypeBase shape o
Mem Space
s

rephraseBinOp ::
  (Monad f) =>
  Rephraser f from rep ->
  SegBinOp from ->
  f (SegBinOp rep)
rephraseBinOp :: forall (f :: * -> *) from rep.
Monad f =>
Rephraser f from rep -> SegBinOp from -> f (SegBinOp rep)
rephraseBinOp Rephraser f from rep
r (SegBinOp Commutativity
comm Lambda from
lam [SubExp]
nes Shape
shape) =
  Commutativity -> Lambda rep -> [SubExp] -> Shape -> SegBinOp rep
forall rep.
Commutativity -> Lambda rep -> [SubExp] -> Shape -> SegBinOp rep
SegBinOp Commutativity
comm (Lambda rep -> [SubExp] -> Shape -> SegBinOp rep)
-> f (Lambda rep) -> f ([SubExp] -> Shape -> SegBinOp rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser f from rep -> Lambda from -> f (Lambda rep)
forall (m :: * -> *) from to t.
Monad m =>
Rephraser m from to -> GLambda from t -> m (GLambda to t)
rephraseLambda Rephraser f from rep
r Lambda from
lam f ([SubExp] -> Shape -> SegBinOp rep)
-> f [SubExp] -> f (Shape -> SegBinOp rep)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [SubExp] -> f [SubExp]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
nes f (Shape -> SegBinOp rep) -> f Shape -> f (SegBinOp rep)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Shape -> f Shape
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Shape
shape

rephrasePostOp ::
  (Monad f) =>
  Rephraser f from rep ->
  SegPostOp from ->
  f (SegPostOp rep)
rephrasePostOp :: forall (f :: * -> *) from rep.
Monad f =>
Rephraser f from rep -> SegPostOp from -> f (SegPostOp rep)
rephrasePostOp Rephraser f from rep
r (SegPostOp Lambda from
lam) =
  Lambda rep -> SegPostOp rep
forall rep. Lambda rep -> SegPostOp rep
SegPostOp (Lambda rep -> SegPostOp rep)
-> f (Lambda rep) -> f (SegPostOp rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser f from rep -> Lambda from -> f (Lambda rep)
forall (m :: * -> *) from to t.
Monad m =>
Rephraser m from to -> GLambda from t -> m (GLambda to t)
rephraseLambda Rephraser f from rep
r Lambda from
lam

instance RephraseOp (SegOp lvl) where
  rephraseInOp :: forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> SegOp lvl from -> m (SegOp lvl to)
rephraseInOp Rephraser m from to
r (SegMap lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody from
body) =
    lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody to
-> SegOp lvl to
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> SegOp lvl rep
SegMap lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts (KernelBody to -> SegOp lvl to)
-> m (KernelBody to) -> m (SegOp lvl to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser m from to -> KernelBody from -> m (KernelBody to)
forall (m :: * -> *) from to res.
Monad m =>
Rephraser m from to -> GBody from res -> m (GBody to res)
rephraseBody Rephraser m from to
r KernelBody from
body
  rephraseInOp Rephraser m from to
r (SegRed lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody from
body [SegBinOp from]
reds) =
    lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody to
-> [SegBinOp to]
-> SegOp lvl to
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
SegRed lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts
      (KernelBody to -> [SegBinOp to] -> SegOp lvl to)
-> m (KernelBody to) -> m ([SegBinOp to] -> SegOp lvl to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser m from to -> KernelBody from -> m (KernelBody to)
forall (m :: * -> *) from to res.
Monad m =>
Rephraser m from to -> GBody from res -> m (GBody to res)
rephraseBody Rephraser m from to
r KernelBody from
body
      m ([SegBinOp to] -> SegOp lvl to)
-> m [SegBinOp to] -> m (SegOp lvl to)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SegBinOp from -> m (SegBinOp to))
-> [SegBinOp from] -> m [SegBinOp to]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Rephraser m from to -> SegBinOp from -> m (SegBinOp to)
forall (f :: * -> *) from rep.
Monad f =>
Rephraser f from rep -> SegBinOp from -> f (SegBinOp rep)
rephraseBinOp Rephraser m from to
r) [SegBinOp from]
reds
  rephraseInOp Rephraser m from to
r (SegScan lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody from
body [SegBinOp from]
scans SegPostOp from
post_op) =
    lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody to
-> [SegBinOp to]
-> SegPostOp to
-> SegOp lvl to
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegPostOp rep
-> SegOp lvl rep
SegScan lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts
      (KernelBody to -> [SegBinOp to] -> SegPostOp to -> SegOp lvl to)
-> m (KernelBody to)
-> m ([SegBinOp to] -> SegPostOp to -> SegOp lvl to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser m from to -> KernelBody from -> m (KernelBody to)
forall (m :: * -> *) from to res.
Monad m =>
Rephraser m from to -> GBody from res -> m (GBody to res)
rephraseBody Rephraser m from to
r KernelBody from
body
      m ([SegBinOp to] -> SegPostOp to -> SegOp lvl to)
-> m [SegBinOp to] -> m (SegPostOp to -> SegOp lvl to)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SegBinOp from -> m (SegBinOp to))
-> [SegBinOp from] -> m [SegBinOp to]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Rephraser m from to -> SegBinOp from -> m (SegBinOp to)
forall (f :: * -> *) from rep.
Monad f =>
Rephraser f from rep -> SegBinOp from -> f (SegBinOp rep)
rephraseBinOp Rephraser m from to
r) [SegBinOp from]
scans
      m (SegPostOp to -> SegOp lvl to)
-> m (SegPostOp to) -> m (SegOp lvl to)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Rephraser m from to -> SegPostOp from -> m (SegPostOp to)
forall (f :: * -> *) from rep.
Monad f =>
Rephraser f from rep -> SegPostOp from -> f (SegPostOp rep)
rephrasePostOp Rephraser m from to
r SegPostOp from
post_op
  rephraseInOp Rephraser m from to
r (SegHist lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody from
body [HistOp from]
hists) =
    lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody to
-> [HistOp to]
-> SegOp lvl to
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [HistOp rep]
-> SegOp lvl rep
SegHist lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts
      (KernelBody to -> [HistOp to] -> SegOp lvl to)
-> m (KernelBody to) -> m ([HistOp to] -> SegOp lvl to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser m from to -> KernelBody from -> m (KernelBody to)
forall (m :: * -> *) from to res.
Monad m =>
Rephraser m from to -> GBody from res -> m (GBody to res)
rephraseBody Rephraser m from to
r KernelBody from
body
      m ([HistOp to] -> SegOp lvl to)
-> m [HistOp to] -> m (SegOp lvl to)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (HistOp from -> m (HistOp to)) -> [HistOp from] -> m [HistOp to]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM HistOp from -> m (HistOp to)
onOp [HistOp from]
hists
    where
      onOp :: HistOp from -> m (HistOp to)
onOp (HistOp Shape
w SubExp
rf [VName]
arrs [SubExp]
nes Shape
shape Lambda from
op) =
        Shape
-> SubExp -> [VName] -> [SubExp] -> Shape -> Lambda to -> HistOp to
forall rep.
Shape
-> SubExp
-> [VName]
-> [SubExp]
-> Shape
-> Lambda rep
-> HistOp rep
HistOp Shape
w SubExp
rf [VName]
arrs [SubExp]
nes Shape
shape (Lambda to -> HistOp to) -> m (Lambda to) -> m (HistOp to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser m from to -> Lambda from -> m (Lambda to)
forall (m :: * -> *) from to t.
Monad m =>
Rephraser m from to -> GLambda from t -> m (GLambda to t)
rephraseLambda Rephraser m from to
r Lambda from
op

-- | A helper for defining 'TraverseOpStms'.
traverseSegOpStms :: (Monad m) => OpStmsTraverser m (SegOp lvl rep) rep
traverseSegOpStms :: forall (m :: * -> *) lvl rep.
Monad m =>
OpStmsTraverser m (SegOp lvl rep) rep
traverseSegOpStms Scope rep -> Stms rep -> m (Stms rep)
f SegOp lvl rep
segop = SegOpMapper lvl rep rep m -> SegOp lvl rep -> m (SegOp lvl rep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper lvl rep rep m
mapper SegOp lvl rep
segop
  where
    seg_scope :: Scope rep
seg_scope = SegSpace -> Scope rep
forall rep. SegSpace -> Scope rep
scopeOfSegSpace (SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
segop)
    f' :: Scope rep -> Stms rep -> m (Stms rep)
f' Scope rep
scope = Scope rep -> Stms rep -> m (Stms rep)
f (Scope rep
seg_scope Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> Scope rep
scope)
    mapper :: SegOpMapper lvl rep rep m
mapper =
      SegOpMapper lvl rep rep m
forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
identitySegOpMapper
        { mapOnSegBinOpLambda = traverseLambdaStms f',
          mapOnSegOpBody = onBody
        }
    onBody :: GBody rep KernelResult -> m (GBody rep KernelResult)
onBody (Body BodyDec rep
dec Stms rep
stms [KernelResult]
res) =
      BodyDec rep -> Stms rep -> [KernelResult] -> GBody rep KernelResult
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body BodyDec rep
dec (Stms rep -> [KernelResult] -> GBody rep KernelResult)
-> m (Stms rep) -> m ([KernelResult] -> GBody rep KernelResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Scope rep -> Stms rep -> m (Stms rep)
f Scope rep
seg_scope Stms rep
stms m ([KernelResult] -> GBody rep KernelResult)
-> m [KernelResult] -> m (GBody rep KernelResult)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [KernelResult] -> m [KernelResult]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [KernelResult]
res

instance
  (ASTRep rep, Substitute lvl) =>
  Substitute (SegOp lvl rep)
  where
  substituteNames :: Map VName VName -> SegOp lvl rep -> SegOp lvl rep
substituteNames Map VName VName
subst = Identity (SegOp lvl rep) -> SegOp lvl rep
forall a. Identity a -> a
runIdentity (Identity (SegOp lvl rep) -> SegOp lvl rep)
-> (SegOp lvl rep -> Identity (SegOp lvl rep))
-> SegOp lvl rep
-> SegOp lvl rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegOpMapper lvl rep rep Identity
-> SegOp lvl rep -> Identity (SegOp lvl rep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper lvl rep rep Identity
substitute
    where
      substitute :: SegOpMapper lvl rep rep Identity
substitute =
        SegOpMapper
          { mapOnSegOpSubExp :: SubExp -> Identity SubExp
mapOnSegOpSubExp = SubExp -> Identity SubExp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> Identity SubExp)
-> (SubExp -> SubExp) -> SubExp -> Identity SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Map VName VName -> SubExp -> SubExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst,
            mapOnSegBinOpLambda :: Lambda rep -> Identity (Lambda rep)
mapOnSegBinOpLambda = Lambda rep -> Identity (Lambda rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda rep -> Identity (Lambda rep))
-> (Lambda rep -> Lambda rep)
-> Lambda rep
-> Identity (Lambda rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Map VName VName -> Lambda rep -> Lambda rep
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst,
            mapOnSegPostOpLambda :: Lambda rep -> Identity (Lambda rep)
mapOnSegPostOpLambda = Lambda rep -> Identity (Lambda rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda rep -> Identity (Lambda rep))
-> (Lambda rep -> Lambda rep)
-> Lambda rep
-> Identity (Lambda rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Map VName VName -> Lambda rep -> Lambda rep
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst,
            mapOnSegOpBody :: KernelBody rep -> Identity (KernelBody rep)
mapOnSegOpBody = KernelBody rep -> Identity (KernelBody rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (KernelBody rep -> Identity (KernelBody rep))
-> (KernelBody rep -> KernelBody rep)
-> KernelBody rep
-> Identity (KernelBody rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Map VName VName -> KernelBody rep -> KernelBody rep
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst,
            mapOnSegOpVName :: VName -> Identity VName
mapOnSegOpVName = VName -> Identity VName
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Identity VName)
-> (VName -> VName) -> VName -> Identity VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst,
            mapOnSegOpLevel :: lvl -> Identity lvl
mapOnSegOpLevel = lvl -> Identity lvl
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (lvl -> Identity lvl) -> (lvl -> lvl) -> lvl -> Identity lvl
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Map VName VName -> lvl -> lvl
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst
          }

instance (ASTRep rep, ASTConstraints lvl) => Rename (SegOp lvl rep) where
  rename :: SegOp lvl rep -> RenameM (SegOp lvl rep)
rename SegOp lvl rep
op =
    [VName] -> RenameM (SegOp lvl rep) -> RenameM (SegOp lvl rep)
forall a. [VName] -> RenameM a -> RenameM a
renameBound (Map VName (NameInfo (ZonkAny 3)) -> [VName]
forall k a. Map k a -> [k]
M.keys (SegSpace -> Map VName (NameInfo (ZonkAny 3))
forall rep. SegSpace -> Scope rep
scopeOfSegSpace (SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op))) (RenameM (SegOp lvl rep) -> RenameM (SegOp lvl rep))
-> RenameM (SegOp lvl rep) -> RenameM (SegOp lvl rep)
forall a b. (a -> b) -> a -> b
$ SegOpMapper lvl rep rep RenameM
-> SegOp lvl rep -> RenameM (SegOp lvl rep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper lvl rep rep RenameM
renamer SegOp lvl rep
op
    where
      renamer :: SegOpMapper lvl rep rep RenameM
renamer = (SubExp -> RenameM SubExp)
-> (Lambda rep -> RenameM (Lambda rep))
-> (Lambda rep -> RenameM (Lambda rep))
-> (KernelBody rep -> RenameM (KernelBody rep))
-> (VName -> RenameM VName)
-> (lvl -> RenameM lvl)
-> SegOpMapper lvl rep rep RenameM
forall lvl frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Lambda frep -> m (Lambda trep))
-> (Lambda frep -> m (Lambda trep))
-> (KernelBody frep -> m (KernelBody trep))
-> (VName -> m VName)
-> (lvl -> m lvl)
-> SegOpMapper lvl frep trep m
SegOpMapper SubExp -> RenameM SubExp
forall a. Rename a => a -> RenameM a
rename Lambda rep -> RenameM (Lambda rep)
forall a. Rename a => a -> RenameM a
rename Lambda rep -> RenameM (Lambda rep)
forall a. Rename a => a -> RenameM a
rename KernelBody rep -> RenameM (KernelBody rep)
forall a. Rename a => a -> RenameM a
rename VName -> RenameM VName
forall a. Rename a => a -> RenameM a
rename lvl -> RenameM lvl
forall a. Rename a => a -> RenameM a
rename

instance (ASTRep rep, FreeIn lvl) => FreeIn (SegOp lvl rep) where
  freeIn' :: SegOp lvl rep -> FV
freeIn' SegOp lvl rep
e =
    Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Map VName (NameInfo (ZonkAny 2)) -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName (NameInfo (ZonkAny 2)) -> [VName])
-> Map VName (NameInfo (ZonkAny 2)) -> [VName]
forall a b. (a -> b) -> a -> b
$ SegSpace -> Map VName (NameInfo (ZonkAny 2))
forall rep. SegSpace -> Scope rep
scopeOfSegSpace (SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
e)) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$
      (State FV (SegOp lvl rep) -> FV -> FV)
-> FV -> State FV (SegOp lvl rep) -> FV
forall a b c. (a -> b -> c) -> b -> a -> c
flip State FV (SegOp lvl rep) -> FV -> FV
forall s a. State s a -> s -> s
execState FV
forall a. Monoid a => a
mempty (State FV (SegOp lvl rep) -> FV) -> State FV (SegOp lvl rep) -> FV
forall a b. (a -> b) -> a -> b
$
        SegOpMapper lvl rep rep (StateT FV Identity)
-> SegOp lvl rep -> State FV (SegOp lvl rep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper lvl rep rep (StateT FV Identity)
free SegOp lvl rep
e
    where
      walk :: (b -> s) -> b -> m b
walk b -> s
f b
x = (s -> s) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (s -> s -> s
forall a. Semigroup a => a -> a -> a
<> b -> s
f b
x) m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
x
      free :: SegOpMapper lvl rep rep (StateT FV Identity)
free =
        SegOpMapper
          { mapOnSegOpSubExp :: SubExp -> StateT FV Identity SubExp
mapOnSegOpSubExp = (SubExp -> FV) -> SubExp -> StateT FV Identity SubExp
forall {m :: * -> *} {s} {b}.
(MonadState s m, Semigroup s) =>
(b -> s) -> b -> m b
walk SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn',
            mapOnSegBinOpLambda :: Lambda rep -> StateT FV Identity (Lambda rep)
mapOnSegBinOpLambda = (Lambda rep -> FV) -> Lambda rep -> StateT FV Identity (Lambda rep)
forall {m :: * -> *} {s} {b}.
(MonadState s m, Semigroup s) =>
(b -> s) -> b -> m b
walk Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
            mapOnSegPostOpLambda :: Lambda rep -> StateT FV Identity (Lambda rep)
mapOnSegPostOpLambda = (Lambda rep -> FV) -> Lambda rep -> StateT FV Identity (Lambda rep)
forall {m :: * -> *} {s} {b}.
(MonadState s m, Semigroup s) =>
(b -> s) -> b -> m b
walk Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
            mapOnSegOpBody :: KernelBody rep -> StateT FV Identity (KernelBody rep)
mapOnSegOpBody = (KernelBody rep -> FV)
-> KernelBody rep -> StateT FV Identity (KernelBody rep)
forall {m :: * -> *} {s} {b}.
(MonadState s m, Semigroup s) =>
(b -> s) -> b -> m b
walk KernelBody rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
            mapOnSegOpVName :: VName -> StateT FV Identity VName
mapOnSegOpVName = (VName -> FV) -> VName -> StateT FV Identity VName
forall {m :: * -> *} {s} {b}.
(MonadState s m, Semigroup s) =>
(b -> s) -> b -> m b
walk VName -> FV
forall a. FreeIn a => a -> FV
freeIn',
            mapOnSegOpLevel :: lvl -> StateT FV Identity lvl
mapOnSegOpLevel = (lvl -> FV) -> lvl -> StateT FV Identity lvl
forall {m :: * -> *} {s} {b}.
(MonadState s m, Semigroup s) =>
(b -> s) -> b -> m b
walk lvl -> FV
forall a. FreeIn a => a -> FV
freeIn'
          }

instance (OpMetrics (Op rep)) => OpMetrics (SegOp lvl rep) where
  opMetrics :: SegOp lvl rep -> MetricsM ()
opMetrics (SegMap lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"SegMap" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ KernelBody rep -> MetricsM ()
forall rep res. OpMetrics (Op rep) => GBody rep res -> MetricsM ()
bodyMetrics KernelBody rep
body
  opMetrics (SegRed lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body [SegBinOp rep]
reds) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"SegRed" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ do
      (SegBinOp rep -> MetricsM ()) -> [SegBinOp rep] -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Text -> MetricsM () -> MetricsM ()
inside Text
"SegBinOp" (MetricsM () -> MetricsM ())
-> (SegBinOp rep -> MetricsM ()) -> SegBinOp rep -> MetricsM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. GLambda rep (TypeBase Shape NoMode) -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics (GLambda rep (TypeBase Shape NoMode) -> MetricsM ())
-> (SegBinOp rep -> GLambda rep (TypeBase Shape NoMode))
-> SegBinOp rep
-> MetricsM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegBinOp rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda) [SegBinOp rep]
reds
      KernelBody rep -> MetricsM ()
forall rep res. OpMetrics (Op rep) => GBody rep res -> MetricsM ()
bodyMetrics KernelBody rep
body
  opMetrics (SegScan lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body [SegBinOp rep]
scans SegPostOp rep
post_op) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"SegScan" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ do
      (SegBinOp rep -> MetricsM ()) -> [SegBinOp rep] -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Text -> MetricsM () -> MetricsM ()
inside Text
"SegBinOp" (MetricsM () -> MetricsM ())
-> (SegBinOp rep -> MetricsM ()) -> SegBinOp rep -> MetricsM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. GLambda rep (TypeBase Shape NoMode) -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics (GLambda rep (TypeBase Shape NoMode) -> MetricsM ())
-> (SegBinOp rep -> GLambda rep (TypeBase Shape NoMode))
-> SegBinOp rep
-> MetricsM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegBinOp rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda) [SegBinOp rep]
scans
      Text -> MetricsM () -> MetricsM ()
inside Text
"SegBinOp" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ GLambda rep (TypeBase Shape NoMode) -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics (GLambda rep (TypeBase Shape NoMode) -> MetricsM ())
-> GLambda rep (TypeBase Shape NoMode) -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ SegPostOp rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp rep
post_op
      KernelBody rep -> MetricsM ()
forall rep res. OpMetrics (Op rep) => GBody rep res -> MetricsM ()
bodyMetrics KernelBody rep
body
  opMetrics (SegHist lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
body [HistOp rep]
ops) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"SegHist" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ do
      (HistOp rep -> MetricsM ()) -> [HistOp rep] -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (GLambda rep (TypeBase Shape NoMode) -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics (GLambda rep (TypeBase Shape NoMode) -> MetricsM ())
-> (HistOp rep -> GLambda rep (TypeBase Shape NoMode))
-> HistOp rep
-> MetricsM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. HistOp rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. HistOp rep -> Lambda rep
histOp) [HistOp rep]
ops
      KernelBody rep -> MetricsM ()
forall rep res. OpMetrics (Op rep) => GBody rep res -> MetricsM ()
bodyMetrics KernelBody rep
body

instance Pretty SegSpace where
  pretty :: forall ann. SegSpace -> Doc ann
pretty (SegSpace VName
phys [(VName, SubExp)]
dims) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
apply
      ( do
          (i, d) <- [(VName, SubExp)]
dims
          pure $ pretty i <+> "<" <+> pretty d
      )
      Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann
"~" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
phys)

instance (PrettyRep rep) => Pretty (SegBinOp rep) where
  pretty :: forall ann. SegBinOp rep -> Doc ann
pretty (SegBinOp Commutativity
comm Lambda rep
lam [SubExp]
nes Shape
shape) =
    Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.braces ([Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty [SubExp]
nes)
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.comma
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Shape -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Shape -> Doc ann
pretty Shape
shape
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.comma
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann
comm'
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
lam
    where
      comm' :: Doc ann
comm' = case Commutativity
comm of
        Commutativity
Commutative -> Doc ann
"commutative "
        Commutativity
Noncommutative -> Doc ann
forall a. Monoid a => a
mempty

instance (PrettyRep rep) => Pretty (SegPostOp rep) where
  pretty :: forall ann. SegPostOp rep -> Doc ann
pretty (SegPostOp Lambda rep
lam) =
    Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
lam

instance (PrettyRep rep, PP.Pretty lvl) => PP.Pretty (SegOp lvl rep) where
  pretty :: forall ann. SegOp lvl rep -> Doc ann
pretty (SegMap lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody rep
body) =
    Doc ann
"segmap"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> lvl -> Doc ann
forall ann. lvl -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty lvl
lvl
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.align (SegSpace -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SegSpace -> Doc ann
pretty SegSpace
space)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann
forall ann. Doc ann
PP.colon
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
ppTuple' ((TypeBase Shape NoMode -> Doc ann)
-> [TypeBase Shape NoMode] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase Shape NoMode -> Doc ann
pretty [TypeBase Shape NoMode]
ts)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.nestedBlock (KernelBody rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. KernelBody rep -> Doc ann
pretty KernelBody rep
body)
  pretty (SegRed lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody rep
body [SegBinOp rep]
reds) =
    Doc ann
"segred"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> lvl -> Doc ann
forall ann. lvl -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty lvl
lvl
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.align (SegSpace -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SegSpace -> Doc ann
pretty SegSpace
space)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann
forall ann. Doc ann
PP.colon
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
ppTuple' ((TypeBase Shape NoMode -> Doc ann)
-> [TypeBase Shape NoMode] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase Shape NoMode -> Doc ann
pretty [TypeBase Shape NoMode]
ts)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.nestedBlock (KernelBody rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. KernelBody rep -> Doc ann
pretty KernelBody rep
body)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.parens ([Doc ann] -> Doc ann
forall a. Monoid a => [a] -> a
mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
intersperse (Doc ann
forall ann. Doc ann
PP.comma Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$ (SegBinOp rep -> Doc ann) -> [SegBinOp rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SegBinOp rep -> Doc ann
pretty [SegBinOp rep]
reds)
  pretty (SegScan lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody rep
body [SegBinOp rep]
scans SegPostOp rep
post_op) =
    Doc ann
"segscan"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> lvl -> Doc ann
forall ann. lvl -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty lvl
lvl
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.align (SegSpace -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SegSpace -> Doc ann
pretty SegSpace
space)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann
forall ann. Doc ann
PP.colon
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
ppTuple' ((TypeBase Shape NoMode -> Doc ann)
-> [TypeBase Shape NoMode] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase Shape NoMode -> Doc ann
pretty [TypeBase Shape NoMode]
ts)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.nestedBlock (KernelBody rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. KernelBody rep -> Doc ann
pretty KernelBody rep
body)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.parens ([Doc ann] -> Doc ann
forall a. Monoid a => [a] -> a
mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
intersperse (Doc ann
forall ann. Doc ann
PP.comma Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$ (SegBinOp rep -> Doc ann) -> [SegBinOp rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SegBinOp rep -> Doc ann
pretty [SegBinOp rep]
scans)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> SegPostOp rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SegPostOp rep -> Doc ann
pretty SegPostOp rep
post_op
  pretty (SegHist lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody rep
body [HistOp rep]
ops) =
    Doc ann
"seghist"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> lvl -> Doc ann
forall ann. lvl -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty lvl
lvl
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.align (SegSpace -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SegSpace -> Doc ann
pretty SegSpace
space)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann
forall ann. Doc ann
PP.colon
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
ppTuple' ((TypeBase Shape NoMode -> Doc ann)
-> [TypeBase Shape NoMode] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase Shape NoMode -> Doc ann
pretty [TypeBase Shape NoMode]
ts)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.nestedBlock (KernelBody rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. KernelBody rep -> Doc ann
pretty KernelBody rep
body)
        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.parens ([Doc ann] -> Doc ann
forall a. Monoid a => [a] -> a
mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
intersperse (Doc ann
forall ann. Doc ann
PP.comma Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$ (HistOp rep -> Doc ann) -> [HistOp rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map HistOp rep -> Doc ann
forall {rep} {ann}. PrettyRep rep => HistOp rep -> Doc ann
ppOp [HistOp rep]
ops)
    where
      ppOp :: HistOp rep -> Doc ann
ppOp (HistOp Shape
w SubExp
rf [VName]
dests [SubExp]
nes Shape
shape Lambda rep
op) =
        Shape -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Shape -> Doc ann
pretty Shape
w
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.comma
            Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
rf
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.comma
            Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.braces ([Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ (VName -> Doc ann) -> [VName] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty [VName]
dests)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.comma
            Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.braces ([Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty [SubExp]
nes)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.comma
            Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Shape -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Shape -> Doc ann
pretty Shape
shape
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
PP.comma
            Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
op

instance CanBeAliased (SegOp lvl) where
  addOpAliases :: forall rep.
AliasableRep rep =>
AliasTable -> SegOp lvl rep -> SegOp lvl (Aliases rep)
addOpAliases AliasTable
aliases = Identity (SegOp lvl (Aliases rep)) -> SegOp lvl (Aliases rep)
forall a. Identity a -> a
runIdentity (Identity (SegOp lvl (Aliases rep)) -> SegOp lvl (Aliases rep))
-> (SegOp lvl rep -> Identity (SegOp lvl (Aliases rep)))
-> SegOp lvl rep
-> SegOp lvl (Aliases rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegOpMapper lvl rep (Aliases rep) Identity
-> SegOp lvl rep -> Identity (SegOp lvl (Aliases rep))
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper lvl rep (Aliases rep) Identity
alias
    where
      alias :: SegOpMapper lvl rep (Aliases rep) Identity
alias =
        (SubExp -> Identity SubExp)
-> (Lambda rep -> Identity (Lambda (Aliases rep)))
-> (Lambda rep -> Identity (Lambda (Aliases rep)))
-> (KernelBody rep -> Identity (KernelBody (Aliases rep)))
-> (VName -> Identity VName)
-> (lvl -> Identity lvl)
-> SegOpMapper lvl rep (Aliases rep) Identity
forall lvl frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Lambda frep -> m (Lambda trep))
-> (Lambda frep -> m (Lambda trep))
-> (KernelBody frep -> m (KernelBody trep))
-> (VName -> m VName)
-> (lvl -> m lvl)
-> SegOpMapper lvl frep trep m
SegOpMapper
          SubExp -> Identity SubExp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          (Lambda (Aliases rep) -> Identity (Lambda (Aliases rep))
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda (Aliases rep) -> Identity (Lambda (Aliases rep)))
-> (Lambda rep -> Lambda (Aliases rep))
-> Lambda rep
-> Identity (Lambda (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases)
          (Lambda (Aliases rep) -> Identity (Lambda (Aliases rep))
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda (Aliases rep) -> Identity (Lambda (Aliases rep)))
-> (Lambda rep -> Lambda (Aliases rep))
-> Lambda rep
-> Identity (Lambda (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases)
          (KernelBody (Aliases rep) -> Identity (KernelBody (Aliases rep))
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (KernelBody (Aliases rep) -> Identity (KernelBody (Aliases rep)))
-> (KernelBody rep -> KernelBody (Aliases rep))
-> KernelBody rep
-> Identity (KernelBody (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. AliasTable -> KernelBody rep -> KernelBody (Aliases rep)
forall rep res.
(AliasableRep rep, IsResult res) =>
AliasTable -> GBody rep res -> GBody (Aliases rep) res
Alias.analyseBody AliasTable
aliases)
          VName -> Identity VName
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          lvl -> Identity lvl
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance CanBeWise (SegOp lvl) where
  addOpWisdom :: forall rep. Informing rep => SegOp lvl rep -> SegOp lvl (Wise rep)
addOpWisdom = Identity (SegOp lvl (Wise rep)) -> SegOp lvl (Wise rep)
forall a. Identity a -> a
runIdentity (Identity (SegOp lvl (Wise rep)) -> SegOp lvl (Wise rep))
-> (SegOp lvl rep -> Identity (SegOp lvl (Wise rep)))
-> SegOp lvl rep
-> SegOp lvl (Wise rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegOpMapper lvl rep (Wise rep) Identity
-> SegOp lvl rep -> Identity (SegOp lvl (Wise rep))
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper lvl rep (Wise rep) Identity
forall {lvl}. SegOpMapper lvl rep (Wise rep) Identity
add
    where
      add :: SegOpMapper lvl rep (Wise rep) Identity
add =
        (SubExp -> Identity SubExp)
-> (Lambda rep -> Identity (Lambda (Wise rep)))
-> (Lambda rep -> Identity (Lambda (Wise rep)))
-> (KernelBody rep -> Identity (KernelBody (Wise rep)))
-> (VName -> Identity VName)
-> (lvl -> Identity lvl)
-> SegOpMapper lvl rep (Wise rep) Identity
forall lvl frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Lambda frep -> m (Lambda trep))
-> (Lambda frep -> m (Lambda trep))
-> (KernelBody frep -> m (KernelBody trep))
-> (VName -> m VName)
-> (lvl -> m lvl)
-> SegOpMapper lvl frep trep m
SegOpMapper SubExp -> Identity SubExp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda (Wise rep) -> Identity (Lambda (Wise rep))
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda (Wise rep) -> Identity (Lambda (Wise rep)))
-> (Lambda rep -> Lambda (Wise rep))
-> Lambda rep
-> Identity (Lambda (Wise rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Lambda rep -> Lambda (Wise rep)
forall rep t.
Informing rep =>
GLambda rep t -> GLambda (Wise rep) t
informLambda) (Lambda (Wise rep) -> Identity (Lambda (Wise rep))
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda (Wise rep) -> Identity (Lambda (Wise rep)))
-> (Lambda rep -> Lambda (Wise rep))
-> Lambda rep
-> Identity (Lambda (Wise rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Lambda rep -> Lambda (Wise rep)
forall rep t.
Informing rep =>
GLambda rep t -> GLambda (Wise rep) t
informLambda) (KernelBody (Wise rep) -> Identity (KernelBody (Wise rep))
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (KernelBody (Wise rep) -> Identity (KernelBody (Wise rep)))
-> (KernelBody rep -> KernelBody (Wise rep))
-> KernelBody rep
-> Identity (KernelBody (Wise rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. KernelBody rep -> KernelBody (Wise rep)
forall rep res.
(Informing rep, IsResult res) =>
GBody rep res -> GBody (Wise rep) res
informBody) VName -> Identity VName
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure lvl -> Identity lvl
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | The element at the given index of an array literal with the given prim-typed
-- elements, as a chain of conditionals. The index is assumed to be in bounds,
-- so the last element is selected without testing for it.
indexArrayLit ::
  PrimType -> TPrimExp Int64 VName -> [PrimExp VName] -> Maybe (PrimExp VName)
indexArrayLit :: PrimType
-> TPrimExp Int64 VName -> [PrimExp VName] -> Maybe (PrimExp VName)
indexArrayLit PrimType
t TPrimExp Int64 VName
i = Integer -> [PrimExp VName] -> Maybe (PrimExp VName)
select Integer
0
  where
    select :: Integer -> [PrimExp VName] -> Maybe (PrimExp VName)
select Integer
_ [] = Maybe (PrimExp VName)
forall a. Maybe a
Nothing
    select Integer
_ [PrimExp VName
x] = PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just PrimExp VName
x
    select Integer
j (PrimExp VName
x : [PrimExp VName]
xs) = do
      y <- Integer -> [PrimExp VName] -> Maybe (PrimExp VName)
select (Integer
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) [PrimExp VName]
xs
      pure $ FunExp (condFun t) [untyped (i .==. fromInteger j), x, y] t

instance (ASTRep rep) => ST.IndexOp (SegOp lvl rep) where
  indexOp :: forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int -> SegOp lvl rep -> [TPrimExp Int64 VName] -> Maybe Indexed
indexOp SymbolTable rep
vtable Int
k (SegMap lvl
_ SegSpace
space [TypeBase Shape NoMode]
_ KernelBody rep
kbody) [TPrimExp Int64 VName]
is = do
    Returns ResultMaySimplify _ se <- Int -> [KernelResult] -> Maybe KernelResult
forall int a. Integral int => int -> [a] -> Maybe a
maybeNth Int
k ([KernelResult] -> Maybe KernelResult)
-> [KernelResult] -> Maybe KernelResult
forall a b. (a -> b) -> a -> b
$ KernelBody rep -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody rep
kbody
    guard $ length gtids <= length is
    let idx_table = [(VName, Indexed)] -> Map VName Indexed
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Indexed)] -> Map VName Indexed)
-> [(VName, Indexed)] -> Map VName Indexed
forall a b. (a -> b) -> a -> b
$ [VName] -> [Indexed] -> [(VName, Indexed)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
gtids ([Indexed] -> [(VName, Indexed)])
-> [Indexed] -> [(VName, Indexed)]
forall a b. (a -> b) -> a -> b
$ (TPrimExp Int64 VName -> Indexed)
-> [TPrimExp Int64 VName] -> [Indexed]
forall a b. (a -> b) -> [a] -> [b]
map (Certs -> PrimExp VName -> Indexed
ST.Indexed Certs
forall a. Monoid a => a
mempty (PrimExp VName -> Indexed)
-> (TPrimExp Int64 VName -> PrimExp VName)
-> TPrimExp Int64 VName
-> Indexed
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. TPrimExp Int64 VName -> PrimExp VName
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped) [TPrimExp Int64 VName]
is
    case se of
      Var VName
v ->
        VName -> Map VName Indexed -> Maybe Indexed
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName Indexed -> Maybe Indexed)
-> Map VName Indexed -> Maybe Indexed
forall a b. (a -> b) -> a -> b
$
          (Map VName Indexed -> Stm rep -> Map VName Indexed)
-> Map VName Indexed -> [Stm rep] -> Map VName Indexed
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 Indexed -> Stm rep -> Map VName Indexed
expandIndexedTable Map VName Indexed
idx_table ([Stm rep] -> Map VName Indexed) -> [Stm rep] -> Map VName Indexed
forall a b. (a -> b) -> a -> b
$
            VName -> Seq (Stm rep) -> [Stm rep]
forall {t :: * -> *} {rep}.
(Foldable t, FreeDec (ExpDec rep), FreeDec (BodyDec rep),
 FreeIn (OpC rep rep), FreeIn (FParamInfo rep),
 FreeIn (LParamInfo rep), FreeIn (LetDec rep), FreeIn (RetType rep),
 FreeIn (BranchType rep)) =>
VName -> t (Stm rep) -> [Stm rep]
dependedOn VName
v (KernelBody rep -> Seq (Stm rep)
forall rep res. GBody rep res -> Stms rep
bodyStms KernelBody rep
kbody)
      SubExp
_ -> Maybe Indexed
forall a. Maybe a
Nothing
    where
      ([VName]
gtids, [SubExp]
_) = [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, SubExp)] -> ([VName], [SubExp]))
-> [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
      -- Indexes in excess of what is used to index through the
      -- segment dimensions.
      excess_is :: [TPrimExp Int64 VName]
excess_is = Int -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. Int -> [a] -> [a]
drop ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
gtids) [TPrimExp Int64 VName]
is

      -- The table entry for a variable depends only on the statements
      -- that variable is (transitively) computed from, and computing an
      -- entry is expensive, so skip the rest. See Note [Pruning the
      -- index table].
      dependedOn :: VName -> t (Stm rep) -> [Stm rep]
dependedOn VName
v = (Names, [Stm rep]) -> [Stm rep]
forall a b. (a, b) -> b
snd ((Names, [Stm rep]) -> [Stm rep])
-> (t (Stm rep) -> (Names, [Stm rep])) -> t (Stm rep) -> [Stm rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Stm rep -> (Names, [Stm rep]) -> (Names, [Stm rep]))
-> (Names, [Stm rep]) -> t (Stm rep) -> (Names, [Stm rep])
forall a b. (a -> b -> b) -> b -> t a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Stm rep -> (Names, [Stm rep]) -> (Names, [Stm rep])
forall {rep}.
(FreeDec (ExpDec rep), FreeDec (BodyDec rep), FreeIn (OpC rep rep),
 FreeIn (FParamInfo rep), FreeIn (LParamInfo rep),
 FreeIn (LetDec rep), FreeIn (RetType rep),
 FreeIn (BranchType rep)) =>
Stm rep -> (Names, [Stm rep]) -> (Names, [Stm rep])
onStm (VName -> Names
oneName VName
v, [])
        where
          onStm :: Stm rep -> (Names, [Stm rep]) -> (Names, [Stm rep])
onStm Stm rep
stm (Names
needed, [Stm rep]
acc)
            | (PatElem (LetDec rep) -> Bool) -> [PatElem (LetDec rep)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> Names -> Bool
`nameIn` Names
needed) (VName -> Bool)
-> (PatElem (LetDec rep) -> VName) -> PatElem (LetDec rep) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName) (Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems (Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm)) =
                (Stm rep -> Names
forall a. FreeIn a => a -> Names
freeIn Stm rep
stm Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
needed, Stm rep
stm Stm rep -> [Stm rep] -> [Stm rep]
forall a. a -> [a] -> [a]
: [Stm rep]
acc)
            | Bool
otherwise = (Names
needed, [Stm rep]
acc)

      expandIndexedTable :: Map VName Indexed -> Stm rep -> Map VName Indexed
expandIndexedTable Map VName Indexed
table Stm rep
stm
        | [VName
v] <- Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (LetDec rep) -> [VName]) -> Pat (LetDec rep) -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm,
          Just (PrimExp VName
pe, Certs
cs) <-
            WriterT Certs Maybe (PrimExp VName) -> Maybe (PrimExp VName, Certs)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT Certs Maybe (PrimExp VName)
 -> Maybe (PrimExp VName, Certs))
-> WriterT Certs Maybe (PrimExp VName)
-> Maybe (PrimExp VName, Certs)
forall a b. (a -> b) -> a -> b
$ (VName -> WriterT Certs Maybe (PrimExp VName))
-> Exp rep -> WriterT Certs Maybe (PrimExp VName)
forall (m :: * -> *) rep v.
(MonadFail m, RepTypes rep) =>
(VName -> m (PrimExp v)) -> Exp rep -> m (PrimExp v)
primExpFromExp (Map VName Indexed -> VName -> WriterT Certs Maybe (PrimExp VName)
asPrimExp Map VName Indexed
table) (Exp rep -> WriterT Certs Maybe (PrimExp VName))
-> Exp rep -> WriterT Certs Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
            VName -> Indexed -> Map VName Indexed -> Map VName Indexed
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (Certs -> PrimExp VName -> Indexed
ST.Indexed (Stm rep -> Certs
forall rep. Stm rep -> Certs
stmCerts Stm rep
stm Certs -> Certs -> Certs
forall a. Semigroup a => a -> a -> a
<> Certs
cs) PrimExp VName
pe) Map VName Indexed
table
        | [VName
v] <- Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (LetDec rep) -> [VName]) -> Pat (LetDec rep) -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm,
          BasicOp (Index VName
arr Slice SubExp
slice) <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm,
          [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [TPrimExp Int64 VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TPrimExp Int64 VName]
excess_is,
          VName
arr VName -> SymbolTable rep -> Bool
forall rep. VName -> SymbolTable rep -> Bool
`ST.available` SymbolTable rep
vtable,
          Just (Slice (PrimExp VName)
slice', Certs
cs) <- Map VName Indexed
-> Slice SubExp -> Maybe (Slice (PrimExp VName), Certs)
asPrimExpSlice Map VName Indexed
table Slice SubExp
slice =
            let idx :: Indexed
idx =
                  Certs -> VName -> [TPrimExp Int64 VName] -> Indexed
ST.IndexedArray
                    (Stm rep -> Certs
forall rep. Stm rep -> Certs
stmCerts Stm rep
stm Certs -> Certs -> Certs
forall a. Semigroup a => a -> a -> a
<> Certs
cs)
                    VName
arr
                    (Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice ((PrimExp VName -> TPrimExp Int64 VName)
-> Slice (PrimExp VName) -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PrimExp VName -> TPrimExp Int64 VName
forall v. PrimExp v -> TPrimExp Int64 v
isInt64 Slice (PrimExp VName)
slice') [TPrimExp Int64 VName]
excess_is)
             in VName -> Indexed -> Map VName Indexed -> Map VName Indexed
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v Indexed
idx Map VName Indexed
table
        | [VName
v] <- Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (LetDec rep) -> [VName]) -> Pat (LetDec rep) -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm,
          BasicOp (ArrayLit [SubExp]
ses (Prim PrimType
t)) <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm,
          [TPrimExp Int64 VName
i] <- [TPrimExp Int64 VName]
excess_is,
          Just ([PrimExp VName]
pes, Certs
cs) <- WriterT Certs Maybe [PrimExp VName]
-> Maybe ([PrimExp VName], Certs)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT Certs Maybe [PrimExp VName]
 -> Maybe ([PrimExp VName], Certs))
-> WriterT Certs Maybe [PrimExp VName]
-> Maybe ([PrimExp VName], Certs)
forall a b. (a -> b) -> a -> b
$ (SubExp -> WriterT Certs Maybe (PrimExp VName))
-> [SubExp] -> WriterT Certs Maybe [PrimExp VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((VName -> WriterT Certs Maybe (PrimExp VName))
-> SubExp -> WriterT Certs Maybe (PrimExp VName)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM (Map VName Indexed -> VName -> WriterT Certs Maybe (PrimExp VName)
asPrimExp Map VName Indexed
table)) [SubExp]
ses,
          Just PrimExp VName
pe <- PrimType
-> TPrimExp Int64 VName -> [PrimExp VName] -> Maybe (PrimExp VName)
indexArrayLit PrimType
t TPrimExp Int64 VName
i [PrimExp VName]
pes =
            VName -> Indexed -> Map VName Indexed -> Map VName Indexed
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (Certs -> PrimExp VName -> Indexed
ST.Indexed (Stm rep -> Certs
forall rep. Stm rep -> Certs
stmCerts Stm rep
stm Certs -> Certs -> Certs
forall a. Semigroup a => a -> a -> a
<> Certs
cs) PrimExp VName
pe) Map VName Indexed
table
        | Bool
otherwise =
            Map VName Indexed
table

      asPrimExpSlice :: Map VName Indexed
-> Slice SubExp -> Maybe (Slice (PrimExp VName), Certs)
asPrimExpSlice Map VName Indexed
table =
        WriterT Certs Maybe (Slice (PrimExp VName))
-> Maybe (Slice (PrimExp VName), Certs)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT Certs Maybe (Slice (PrimExp VName))
 -> Maybe (Slice (PrimExp VName), Certs))
-> (Slice SubExp -> WriterT Certs Maybe (Slice (PrimExp VName)))
-> Slice SubExp
-> Maybe (Slice (PrimExp VName), Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (SubExp -> WriterT Certs Maybe (PrimExp VName))
-> Slice SubExp -> WriterT Certs Maybe (Slice (PrimExp VName))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Slice a -> f (Slice b)
traverse ((VName -> WriterT Certs Maybe (PrimExp VName))
-> SubExp -> WriterT Certs Maybe (PrimExp VName)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM (Map VName Indexed -> VName -> WriterT Certs Maybe (PrimExp VName)
asPrimExp Map VName Indexed
table))

      asPrimExp :: Map VName Indexed -> VName -> WriterT Certs Maybe (PrimExp VName)
asPrimExp Map VName Indexed
table VName
v
        | Just (ST.Indexed Certs
cs PrimExp VName
e) <- VName -> Map VName Indexed -> Maybe Indexed
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName Indexed
table = Certs -> WriterT Certs Maybe ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell Certs
cs WriterT Certs Maybe ()
-> WriterT Certs Maybe (PrimExp VName)
-> WriterT Certs Maybe (PrimExp VName)
forall a b.
WriterT Certs Maybe a
-> WriterT Certs Maybe b -> WriterT Certs Maybe b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> PrimExp VName -> WriterT Certs Maybe (PrimExp VName)
forall a. a -> WriterT Certs Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimExp VName
e
        | Just (Prim PrimType
pt) <- VName -> SymbolTable rep -> Maybe (TypeBase Shape NoMode)
forall rep.
ASTRep rep =>
VName -> SymbolTable rep -> Maybe (TypeBase Shape NoMode)
ST.lookupType VName
v SymbolTable rep
vtable =
            PrimExp VName -> WriterT Certs Maybe (PrimExp VName)
forall a. a -> WriterT Certs Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimExp VName -> WriterT Certs Maybe (PrimExp VName))
-> PrimExp VName -> WriterT Certs Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
v PrimType
pt
        | Bool
otherwise = Maybe (PrimExp VName) -> WriterT Certs Maybe (PrimExp VName)
forall (m :: * -> *) a. Monad m => m a -> WriterT Certs m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift Maybe (PrimExp VName)
forall a. Maybe a
Nothing
  indexOp SymbolTable rep
_ Int
_ SegOp lvl rep
_ [TPrimExp Int64 VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

instance (ASTConstraints lvl) => IsOp (SegOp lvl) where
  cheapOp :: forall rep. ASTRep rep => SegOp lvl rep -> Bool
cheapOp SegOp lvl rep
_ = Bool
False
  safeOp :: forall rep. ASTRep rep => SegOp lvl rep -> Bool
safeOp SegOp lvl rep
_ = Bool
True
  opDependencies :: forall rep. ASTRep rep => SegOp lvl rep -> [Names]
opDependencies SegOp lvl rep
op = Int -> Names -> [Names]
forall a. Int -> a -> [a]
replicate ([TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SegOp lvl rep -> [TypeBase Shape NoMode]
forall lvl rep. SegOp lvl rep -> [TypeBase Shape NoMode]
segOpType SegOp lvl rep
op)) (SegOp lvl rep -> Names
forall a. FreeIn a => a -> Names
freeIn SegOp lvl rep
op)

--- Simplification

instance Engine.Simplifiable SegSpace where
  simplify :: forall rep. SimplifiableRep rep => SegSpace -> SimpleM rep SegSpace
simplify (SegSpace VName
phys [(VName, SubExp)]
dims) =
    VName -> [(VName, SubExp)] -> SegSpace
SegSpace VName
phys ([(VName, SubExp)] -> SegSpace)
-> SimpleM rep [(VName, SubExp)] -> SimpleM rep SegSpace
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((VName, SubExp) -> SimpleM rep (VName, SubExp))
-> [(VName, SubExp)] -> SimpleM rep [(VName, SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SubExp -> SimpleM rep SubExp)
-> (VName, SubExp) -> SimpleM rep (VName, SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (VName, a) -> f (VName, b)
traverse SubExp -> SimpleM rep SubExp
forall rep. SimplifiableRep rep => SubExp -> SimpleM rep SubExp
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify) [(VName, SubExp)]
dims

instance Engine.Simplifiable KernelResult where
  simplify :: forall rep.
SimplifiableRep rep =>
KernelResult -> SimpleM rep KernelResult
simplify (Returns ResultManifest
manifest Certs
cs SubExp
what) =
    ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
manifest (Certs -> SubExp -> KernelResult)
-> SimpleM rep Certs -> SimpleM rep (SubExp -> KernelResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Certs -> SimpleM rep Certs
forall rep. SimplifiableRep rep => Certs -> SimpleM rep Certs
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify Certs
cs SimpleM rep (SubExp -> KernelResult)
-> SimpleM rep SubExp -> SimpleM rep KernelResult
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SubExp -> SimpleM rep SubExp
forall rep. SimplifiableRep rep => SubExp -> SimpleM rep SubExp
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify SubExp
what
  simplify (TileReturns Certs
cs [(SubExp, SubExp)]
dims VName
what) =
    Certs -> [(SubExp, SubExp)] -> VName -> KernelResult
TileReturns (Certs -> [(SubExp, SubExp)] -> VName -> KernelResult)
-> SimpleM rep Certs
-> SimpleM rep ([(SubExp, SubExp)] -> VName -> KernelResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Certs -> SimpleM rep Certs
forall rep. SimplifiableRep rep => Certs -> SimpleM rep Certs
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify Certs
cs SimpleM rep ([(SubExp, SubExp)] -> VName -> KernelResult)
-> SimpleM rep [(SubExp, SubExp)]
-> SimpleM rep (VName -> KernelResult)
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(SubExp, SubExp)] -> SimpleM rep [(SubExp, SubExp)]
forall rep.
SimplifiableRep rep =>
[(SubExp, SubExp)] -> SimpleM rep [(SubExp, SubExp)]
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify [(SubExp, SubExp)]
dims SimpleM rep (VName -> KernelResult)
-> SimpleM rep VName -> SimpleM rep KernelResult
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> SimpleM rep VName
forall rep. SimplifiableRep rep => VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
what
  simplify (RegTileReturns Certs
cs [(SubExp, SubExp, SubExp)]
dims_n_tiles VName
what) =
    Certs -> [(SubExp, SubExp, SubExp)] -> VName -> KernelResult
RegTileReturns
      (Certs -> [(SubExp, SubExp, SubExp)] -> VName -> KernelResult)
-> SimpleM rep Certs
-> SimpleM
     rep ([(SubExp, SubExp, SubExp)] -> VName -> KernelResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Certs -> SimpleM rep Certs
forall rep. SimplifiableRep rep => Certs -> SimpleM rep Certs
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify Certs
cs
      SimpleM rep ([(SubExp, SubExp, SubExp)] -> VName -> KernelResult)
-> SimpleM rep [(SubExp, SubExp, SubExp)]
-> SimpleM rep (VName -> KernelResult)
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(SubExp, SubExp, SubExp)]
-> SimpleM rep [(SubExp, SubExp, SubExp)]
forall rep.
SimplifiableRep rep =>
[(SubExp, SubExp, SubExp)]
-> SimpleM rep [(SubExp, SubExp, SubExp)]
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify [(SubExp, SubExp, SubExp)]
dims_n_tiles
      SimpleM rep (VName -> KernelResult)
-> SimpleM rep VName -> SimpleM rep KernelResult
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> SimpleM rep VName
forall rep. SimplifiableRep rep => VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
what

segOpBlocker ::
  (Engine.SimplifiableRep rep) =>
  SegSpace -> Engine.SimpleM rep (Engine.BlockPred (Wise rep))
segOpBlocker :: forall rep.
SimplifiableRep rep =>
SegSpace -> SimpleM rep (BlockPred (Wise rep))
segOpBlocker SegSpace
space = do
  par_blocker <- (Env rep -> BlockPred (Wise rep))
-> SimpleM rep (BlockPred (Wise rep))
forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
Engine.asksEngineEnv ((Env rep -> BlockPred (Wise rep))
 -> SimpleM rep (BlockPred (Wise rep)))
-> (Env rep -> BlockPred (Wise rep))
-> SimpleM rep (BlockPred (Wise rep))
forall a b. (a -> b) -> a -> b
$ HoistBlockers rep -> BlockPred (Wise rep)
forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
Engine.blockHoistPar (HoistBlockers rep -> BlockPred (Wise rep))
-> (Env rep -> HoistBlockers rep)
-> Env rep
-> BlockPred (Wise rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Env rep -> HoistBlockers rep
forall {k} (rep :: k). Env rep -> HoistBlockers rep
Engine.envHoistBlockers
  pure $
    Engine.hasFree bound_here
      `Engine.orIf` Engine.isOp
      `Engine.orIf` isSlice
      `Engine.orIf` par_blocker
      `Engine.orIf` Engine.isConsumed
      `Engine.orIf` Engine.isConsuming
      `Engine.orIf` Engine.isDeviceMigrated
  where
    bound_here :: Names
bound_here = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Map VName (NameInfo (ZonkAny 0)) -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName (NameInfo (ZonkAny 0)) -> [VName])
-> Map VName (NameInfo (ZonkAny 0)) -> [VName]
forall a b. (a -> b) -> a -> b
$ SegSpace -> Map VName (NameInfo (ZonkAny 0))
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space
    -- Hoisting out slicing is useless and potentially dangerous, as protecting
    -- it can cause nonuniform parallelism.
    isSlice :: p -> p -> Stm rep -> Bool
isSlice p
_ p
_ (Let Pat (LetDec rep)
_ StmAux (ExpDec rep)
_ (BasicOp (Index VName
_ Slice SubExp
slice))) =
      Slice SubExp -> Shape
forall d. Slice d -> ShapeBase d
sliceShape Slice SubExp
slice Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
/= Shape
forall a. Monoid a => a
mempty
    isSlice p
_ p
_ Stm rep
_ = Bool
False

-- | We are willing to hoist potentially unsafe statements out of segops, but
-- they must be protected by adding a branch on top of them.
protectSegOpHoisted ::
  (Engine.SimplifiableRep rep) =>
  SegSpace ->
  Engine.SimpleM rep (a, b, Stms (Wise rep)) ->
  Engine.SimpleM rep (a, b, Stms (Wise rep))
protectSegOpHoisted :: forall rep a b.
SimplifiableRep rep =>
SegSpace
-> SimpleM rep (a, b, Stms (Wise rep))
-> SimpleM rep (a, b, Stms (Wise rep))
protectSegOpHoisted SegSpace
space SimpleM rep (a, b, Stms (Wise rep))
m = do
  (x, y, stms) <- SimpleM rep (a, b, Stms (Wise rep))
m
  ops <- asks $ Engine.protectHoistedOpS . fst
  stms' <- runBuilder_ $ do
    if not $ all (safeExp . stmExp) stms
      then do
        is_nonempty <- checkIfNonEmpty
        mapM_ (Engine.protectIf ops (not . safeExp) is_nonempty) stms
      else addStms stms
  pure (x, y, stms')
  where
    checkIfNonEmpty :: BuilderT (Wise rep) (State VNameSource) SubExp
checkIfNonEmpty = do
      segop_size <-
        Name
-> Exp (Rep (BuilderT (Wise rep) (State VNameSource)))
-> BuilderT (Wise rep) (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segop_size"
          (Exp (Wise rep) -> BuilderT (Wise rep) (State VNameSource) SubExp)
-> BuilderT (Wise rep) (State VNameSource) (Exp (Wise rep))
-> BuilderT (Wise rep) (State VNameSource) SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp
-> [SubExp]
-> BuilderT
     (Wise rep)
     (State VNameSource)
     (Exp (Rep (BuilderT (Wise rep) (State VNameSource))))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
OverflowWrap) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1) (SegSpace -> [SubExp]
segSpaceDims SegSpace
space)
      letSubExp "segop_nonempty" . BasicOp $
        CmpOp (CmpSlt Int64) (intConst Int64 0) segop_size

simplifyKernelBody ::
  (Engine.SimplifiableRep rep, BodyDec rep ~ ()) =>
  SegSpace ->
  KernelBody (Wise rep) ->
  Engine.SimpleM rep (KernelBody (Wise rep), Stms (Wise rep))
simplifyKernelBody :: forall rep.
(SimplifiableRep rep, BodyDec rep ~ ()) =>
SegSpace
-> KernelBody (Wise rep)
-> SimpleM rep (KernelBody (Wise rep), Stms (Wise rep))
simplifyKernelBody SegSpace
space (Body BodyDec (Wise rep)
_ Stms (Wise rep)
stms [KernelResult]
res) = do
  blocker <- SegSpace -> SimpleM rep (BlockPred (Wise rep))
forall rep.
SimplifiableRep rep =>
SegSpace -> SimpleM rep (BlockPred (Wise rep))
segOpBlocker SegSpace
space

  -- Ensure we do not try to use anything that is consumed in the result.
  (body_res, body_stms, hoisted) <-
    protectSegOpHoisted space
      . Engine.localVtable (segSpaceSymbolTable space)
      . Engine.localVtable (\SymbolTable (Wise rep)
vtable -> SymbolTable (Wise rep)
vtable {ST.simplifyMemory = True})
      . Engine.enterLoop
      $ Engine.blockIf blocker stms
      $ do
        res' <-
          Engine.localVtable (ST.hideCertified $ namesFromList $ M.keys $ scopeOf stms) $
            mapM Engine.simplify res
        pure (res', UT.usages $ freeIn res')

  pure (mkWiseBody () body_stms body_res, hoisted)

simplifyLambda ::
  (Engine.SimplifiableRep rep) =>
  Names ->
  Lambda (Wise rep) ->
  Engine.SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambda :: forall rep.
SimplifiableRep rep =>
Names
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambda Names
bound = SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall rep.
SimplifiableRep rep =>
SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
Engine.blockMigrated (SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
 -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> (Lambda (Wise rep)
    -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Names
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall rep t.
(SimplifiableRep rep, Simplifiable t) =>
Names
-> GLambda (Wise rep) t
-> SimpleM rep (GLambda (Wise rep) t, Stms (Wise rep))
Engine.simplifyLambda Names
bound

segSpaceSymbolTable :: (ASTRep rep) => SegSpace -> ST.SymbolTable rep -> ST.SymbolTable rep
segSpaceSymbolTable :: forall rep.
ASTRep rep =>
SegSpace -> SymbolTable rep -> SymbolTable rep
segSpaceSymbolTable (SegSpace VName
flat [(VName, SubExp)]
gtids_and_dims) SymbolTable rep
st =
  (SymbolTable rep -> (VName, SubExp) -> SymbolTable rep)
-> SymbolTable rep -> [(VName, SubExp)] -> SymbolTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SymbolTable rep -> (VName, SubExp) -> SymbolTable rep
forall {rep}.
ASTRep rep =>
SymbolTable rep -> (VName, SubExp) -> SymbolTable rep
f (VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep
ST.insertFreeVar VName
flat (IntType -> NameInfo rep
forall rep. IntType -> NameInfo rep
IndexName IntType
Int64) SymbolTable rep
st) [(VName, SubExp)]
gtids_and_dims
  where
    f :: SymbolTable rep -> (VName, SubExp) -> SymbolTable rep
f SymbolTable rep
vtable (VName
gtid, SubExp
dim) = VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep
ST.insertLoopVar VName
gtid IntType
Int64 SubExp
dim SymbolTable rep
vtable

simplifySegBinOp ::
  (Engine.SimplifiableRep rep) =>
  VName ->
  SegBinOp (Wise rep) ->
  Engine.SimpleM rep (SegBinOp (Wise rep), Stms (Wise rep))
simplifySegBinOp :: forall rep.
SimplifiableRep rep =>
VName
-> SegBinOp (Wise rep)
-> SimpleM rep (SegBinOp (Wise rep), Stms (Wise rep))
simplifySegBinOp VName
phys_id (SegBinOp Commutativity
comm Lambda (Wise rep)
lam [SubExp]
nes Shape
shape) = do
  (lam', hoisted) <-
    SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep a
Engine.enterLoop
      (SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
 -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> (SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
    -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
Engine.localVtable (\SymbolTable (Wise rep)
vtable -> SymbolTable (Wise rep)
vtable {ST.simplifyMemory = True})
      (SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
 -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall a b. (a -> b) -> a -> b
$ Names
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall rep.
SimplifiableRep rep =>
Names
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambda (VName -> Names
oneName VName
phys_id) Lambda (Wise rep)
lam
  shape' <- Engine.simplify shape
  nes' <- mapM Engine.simplify nes
  pure (SegBinOp comm lam' nes' shape', hoisted)

simplifySegPostOp ::
  (Engine.SimplifiableRep rep) =>
  SegSpace ->
  SegPostOp (Wise rep) ->
  Engine.SimpleM rep (SegPostOp (Wise rep), Stms (Wise rep))
simplifySegPostOp :: forall rep.
SimplifiableRep rep =>
SegSpace
-> SegPostOp (Wise rep)
-> SimpleM rep (SegPostOp (Wise rep), Stms (Wise rep))
simplifySegPostOp SegSpace
space (SegPostOp Lambda (Wise rep)
lam) = do
  (lam', hoisted) <-
    SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep a
Engine.enterLoop
      (SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
 -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> (SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
    -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
Engine.localVtable (\SymbolTable (Wise rep)
vtable -> SymbolTable (Wise rep)
vtable {ST.simplifyMemory = True})
      (SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
 -> SimpleM rep (Lambda (Wise rep), Stms (Wise rep)))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall a b. (a -> b) -> a -> b
$ Names
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
forall rep.
SimplifiableRep rep =>
Names
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambda Names
bound_here Lambda (Wise rep)
lam
  pure (SegPostOp lam', hoisted)
  where
    bound_here :: Names
bound_here = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Map VName (NameInfo (ZonkAny 1)) -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName (NameInfo (ZonkAny 1)) -> [VName])
-> Map VName (NameInfo (ZonkAny 1)) -> [VName]
forall a b. (a -> b) -> a -> b
$ SegSpace -> Map VName (NameInfo (ZonkAny 1))
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space

-- | Simplify the given 'SegOp'.
simplifySegOp ::
  ( Engine.SimplifiableRep rep,
    BodyDec rep ~ (),
    Engine.Simplifiable lvl
  ) =>
  SegOp lvl (Wise rep) ->
  Engine.SimpleM rep (SegOp lvl (Wise rep), Stms (Wise rep))
simplifySegOp :: forall rep lvl.
(SimplifiableRep rep, BodyDec rep ~ (), Simplifiable lvl) =>
SegOp lvl (Wise rep)
-> SimpleM rep (SegOp lvl (Wise rep), Stms (Wise rep))
simplifySegOp (SegMap lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Wise rep)
kbody) = do
  (lvl', space', ts') <- (lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall rep.
SimplifiableRep rep =>
(lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify (lvl
lvl, SegSpace
space, [TypeBase Shape NoMode]
ts)
  (kbody', body_hoisted) <- simplifyKernelBody space kbody
  pure
    ( SegMap lvl' space' ts' kbody',
      body_hoisted
    )
simplifySegOp (SegRed lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Wise rep)
kbody [SegBinOp (Wise rep)]
reds) = do
  (lvl', space', ts') <- (lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall rep.
SimplifiableRep rep =>
(lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify (lvl
lvl, SegSpace
space, [TypeBase Shape NoMode]
ts)
  (reds', reds_hoisted) <-
    Engine.localVtable (ST.insertScope scope) $
      mapAndUnzipM (simplifySegBinOp (segFlat space)) reds
  (kbody', body_hoisted) <- simplifyKernelBody space kbody

  pure
    ( SegRed lvl' space' ts' kbody' reds',
      mconcat reds_hoisted <> body_hoisted
    )
  where
    scope :: Scope (Wise rep)
scope = SegSpace -> Scope (Wise rep)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space
simplifySegOp (SegScan lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Wise rep)
kbody [SegBinOp (Wise rep)]
scans SegPostOp (Wise rep)
post_op) = do
  (lvl', space', ts') <- (lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall rep.
SimplifiableRep rep =>
(lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify (lvl
lvl, SegSpace
space, [TypeBase Shape NoMode]
ts)
  (scans', scans_hoisted) <-
    Engine.localVtable (ST.insertScope scope) $
      mapAndUnzipM (simplifySegBinOp (segFlat space)) scans
  (kbody', body_hoisted) <- simplifyKernelBody space kbody
  (post_op', post_op_hoisted) <-
    Engine.localVtable (ST.insertScope scope) $
      simplifySegPostOp space post_op

  pure
    ( SegScan lvl' space' ts' kbody' scans' post_op',
      mconcat scans_hoisted <> body_hoisted <> post_op_hoisted
    )
  where
    scope :: Scope (Wise rep)
scope = SegSpace -> Scope (Wise rep)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space
simplifySegOp (SegHist lvl
lvl SegSpace
space [TypeBase Shape NoMode]
ts KernelBody (Wise rep)
kbody [HistOp (Wise rep)]
ops) = do
  (lvl', space', ts') <- (lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall rep.
SimplifiableRep rep =>
(lvl, SegSpace, [TypeBase Shape NoMode])
-> SimpleM rep (lvl, SegSpace, [TypeBase Shape NoMode])
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify (lvl
lvl, SegSpace
space, [TypeBase Shape NoMode]
ts)

  Engine.localVtable (flip (foldr ST.consume) $ concatMap histDest ops) $ do
    (ops', ops_hoisted) <- fmap unzip . forM ops $
      \(HistOp Shape
w SubExp
rf [VName]
arrs [SubExp]
nes Shape
dims Lambda (Wise rep)
lam) -> do
        w' <- Shape -> SimpleM rep Shape
forall rep. SimplifiableRep rep => Shape -> SimpleM rep Shape
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify Shape
w
        rf' <- Engine.simplify rf
        arrs' <- Engine.simplify arrs
        nes' <- Engine.simplify nes
        dims' <- Engine.simplify dims
        (lam', op_hoisted) <-
          Engine.localVtable (ST.insertScope scope) $
            Engine.localVtable (\SymbolTable (Wise rep)
vtable -> SymbolTable (Wise rep)
vtable {ST.simplifyMemory = True}) $
              simplifyLambda (oneName (segFlat space)) lam
        pure
          ( HistOp w' rf' arrs' nes' dims' lam',
            op_hoisted
          )

    (kbody', body_hoisted) <- simplifyKernelBody space kbody

    pure
      ( SegHist lvl' space' ts' kbody' ops',
        mconcat ops_hoisted <> body_hoisted
      )
  where
    scope :: Scope (Wise rep)
scope = SegSpace -> Scope (Wise rep)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space

-- | Does this rep contain 'SegOp's in its t'Op's?  A rep must be an
-- instance of this class for the simplification rules to work.
class HasSegOp rep where
  type SegOpLevel rep
  asSegOp :: Op rep -> Maybe (SegOp (SegOpLevel rep) rep)
  segOp :: SegOp (SegOpLevel rep) rep -> Op rep

-- | Simplification rules for simplifying 'SegOp's.
segOpRules ::
  (HasSegOp rep, BuilderOps rep, Buildable rep, Aliased rep) =>
  RuleBook rep
segOpRules :: forall rep.
(HasSegOp rep, BuilderOps rep, Buildable rep, Aliased rep) =>
RuleBook rep
segOpRules =
  [TopDownRule rep] -> [BottomUpRule rep] -> RuleBook rep
forall m. [TopDownRule m] -> [BottomUpRule m] -> RuleBook m
ruleBook [RuleOp rep (TopDown rep) -> TopDownRule rep
forall rep a. RuleOp rep a -> SimplificationRule rep a
RuleOp RuleOp rep (TopDown rep)
forall rep.
(HasSegOp rep, BuilderOps rep, Buildable rep) =>
TopDownRuleOp rep
segOpRuleTopDown] [RuleOp rep (BottomUp rep) -> BottomUpRule rep
forall rep a. RuleOp rep a -> SimplificationRule rep a
RuleOp RuleOp rep (BottomUp rep)
forall rep.
(HasSegOp rep, BuilderOps rep, Aliased rep) =>
BottomUpRuleOp rep
segOpRuleBottomUp]

segOpRuleTopDown ::
  (HasSegOp rep, BuilderOps rep, Buildable rep) =>
  TopDownRuleOp rep
segOpRuleTopDown :: forall rep.
(HasSegOp rep, BuilderOps rep, Buildable rep) =>
TopDownRuleOp rep
segOpRuleTopDown TopDown rep
vtable Pat (LetDec rep)
pat StmAux (ExpDec rep)
dec Op rep
op
  | Just SegOp (SegOpLevel rep) rep
op' <- Op rep -> Maybe (SegOp (SegOpLevel rep) rep)
forall rep.
HasSegOp rep =>
Op rep -> Maybe (SegOp (SegOpLevel rep) rep)
asSegOp Op rep
op =
      TopDown rep
-> Pat (LetDec rep)
-> StmAux (ExpDec rep)
-> SegOp (SegOpLevel rep) rep
-> Rule rep
forall rep.
(HasSegOp rep, BuilderOps rep, Buildable rep) =>
SymbolTable rep
-> Pat (LetDec rep)
-> StmAux (ExpDec rep)
-> SegOp (SegOpLevel rep) rep
-> Rule rep
topDownSegOp TopDown rep
vtable Pat (LetDec rep)
pat StmAux (ExpDec rep)
dec SegOp (SegOpLevel rep) rep
op'
  | Bool
otherwise =
      Rule rep
forall rep. Rule rep
Skip

segOpRuleBottomUp ::
  (HasSegOp rep, BuilderOps rep, Aliased rep) =>
  BottomUpRuleOp rep
segOpRuleBottomUp :: forall rep.
(HasSegOp rep, BuilderOps rep, Aliased rep) =>
BottomUpRuleOp rep
segOpRuleBottomUp BottomUp rep
vtable Pat (LetDec rep)
pat StmAux (ExpDec rep)
dec Op rep
op
  | Just SegOp (SegOpLevel rep) rep
op' <- Op rep -> Maybe (SegOp (SegOpLevel rep) rep)
forall rep.
HasSegOp rep =>
Op rep -> Maybe (SegOp (SegOpLevel rep) rep)
asSegOp Op rep
op =
      BottomUp rep
-> Pat (LetDec rep)
-> StmAux (ExpDec rep)
-> SegOp (SegOpLevel rep) rep
-> Rule rep
forall rep.
(Aliased rep, HasSegOp rep, BuilderOps rep) =>
(SymbolTable rep, UsageTable)
-> Pat (LetDec rep)
-> StmAux (ExpDec rep)
-> SegOp (SegOpLevel rep) rep
-> Rule rep
bottomUpSegOp BottomUp rep
vtable Pat (LetDec rep)
pat StmAux (ExpDec rep)
dec SegOp (SegOpLevel rep) rep
op'
  | Bool
otherwise =
      Rule rep
forall rep. Rule rep
Skip

-- | Only handle Returns cases.
depsOfRes :: Dependencies -> KernelResult -> Names
depsOfRes :: AliasTable -> KernelResult -> Names
depsOfRes AliasTable
deps (Returns ResultManifest
_ Certs
cs SubExp
se) = AliasTable -> SubExp -> Names
depsOf AliasTable
deps SubExp
se Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> AliasTable -> Names -> Names
depsOfNames AliasTable
deps (Certs -> Names
forall a. FreeIn a => a -> Names
freeIn Certs
cs)
depsOfRes AliasTable
_ KernelResult
_ = Names
forall a. Monoid a => a
mempty

kernelBodyDependencies :: (ASTRep rep) => Dependencies -> KernelBody rep -> [Names]
kernelBodyDependencies :: forall rep. ASTRep rep => AliasTable -> KernelBody rep -> [Names]
kernelBodyDependencies AliasTable
deps KernelBody rep
kbody =
  let names_in_scope :: Names
names_in_scope = KernelBody rep -> Names
forall a. FreeIn a => a -> Names
freeIn KernelBody rep
kbody
      deps' :: AliasTable
deps' = AliasTable -> KernelBody rep -> AliasTable
forall rep res.
ASTRep rep =>
AliasTable -> GBody rep res -> AliasTable
dataDependencies' AliasTable
deps KernelBody rep
kbody
   in (KernelResult -> Names) -> [KernelResult] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map
        ((Names -> Names -> Names) -> Names -> Names -> Names
forall a b c. (a -> b -> c) -> b -> a -> c
flip Names -> Names -> Names
namesSubtract Names
names_in_scope (Names -> Names)
-> (KernelResult -> Names) -> KernelResult -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. AliasTable -> KernelResult -> Names
depsOfRes AliasTable
deps')
        (KernelBody rep -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody rep
kbody)

topDownSegOp ::
  (HasSegOp rep, BuilderOps rep, Buildable rep) =>
  ST.SymbolTable rep ->
  Pat (LetDec rep) ->
  StmAux (ExpDec rep) ->
  SegOp (SegOpLevel rep) rep ->
  Rule rep
-- If a SegOp produces something invariant to the SegOp, turn it
-- into a replicate.
topDownSegOp :: forall rep.
(HasSegOp rep, BuilderOps rep, Buildable rep) =>
SymbolTable rep
-> Pat (LetDec rep)
-> StmAux (ExpDec rep)
-> SegOp (SegOpLevel rep) rep
-> Rule rep
topDownSegOp SymbolTable rep
vtable (Pat [PatElem (LetDec rep)]
kpes) StmAux (ExpDec rep)
dec (SegMap SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
ts (Body BodyDec rep
_ Stms rep
kstms [KernelResult]
kres)) = RuleM rep () -> Rule rep
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM rep () -> Rule rep) -> RuleM rep () -> Rule rep
forall a b. (a -> b) -> a -> b
$ do
  (ts', kpes', kres') <-
    [(TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)]
-> ([TypeBase Shape NoMode], [PatElem (LetDec rep)],
    [KernelResult])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)]
 -> ([TypeBase Shape NoMode], [PatElem (LetDec rep)],
     [KernelResult]))
-> RuleM
     rep [(TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)]
-> RuleM
     rep
     ([TypeBase Shape NoMode], [PatElem (LetDec rep)], [KernelResult])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)
 -> RuleM rep Bool)
-> [(TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)]
-> RuleM
     rep [(TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)
-> RuleM rep Bool
checkForInvarianceResult ([TypeBase Shape NoMode]
-> [PatElem (LetDec rep)]
-> [KernelResult]
-> [(TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [TypeBase Shape NoMode]
ts [PatElem (LetDec rep)]
kpes [KernelResult]
kres)

  -- Check if we did anything at all.
  when (kres == kres') cannotSimplify

  kbody <- mkBodyM kstms kres'
  addStm $ Let (Pat kpes') dec $ Op $ segOp $ SegMap lvl space ts' kbody
  where
    isInvariant :: SubExp -> Bool
isInvariant Constant {} = Bool
True
    isInvariant (Var VName
v) = Maybe (Entry rep) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Entry rep) -> Bool) -> Maybe (Entry rep) -> Bool
forall a b. (a -> b) -> a -> b
$ VName -> SymbolTable rep -> Maybe (Entry rep)
forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
ST.lookup VName
v SymbolTable rep
vtable

    checkForInvarianceResult :: (TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)
-> RuleM rep Bool
checkForInvarianceResult (TypeBase Shape NoMode
_, PatElem (LetDec rep)
pe, Returns ResultManifest
rm Certs
cs SubExp
se)
      | Certs
cs Certs -> Certs -> Bool
forall a. Eq a => a -> a -> Bool
== Certs
forall a. Monoid a => a
mempty,
        ResultManifest
rm ResultManifest -> ResultManifest -> Bool
forall a. Eq a => a -> a -> Bool
== ResultManifest
ResultMaySimplify,
        SubExp -> Bool
isInvariant SubExp
se = do
          [VName] -> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe] (Exp (Rep (RuleM rep)) -> RuleM rep ())
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall a b. (a -> b) -> a -> b
$
            BasicOp -> Exp (Rep (RuleM rep))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (RuleM rep)))
-> BasicOp -> Exp (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$
              Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ SegSpace -> [SubExp]
segSpaceDims SegSpace
space) SubExp
se
          Bool -> RuleM rep Bool
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
    checkForInvarianceResult (TypeBase Shape NoMode, PatElem (LetDec rep), KernelResult)
_ =
      Bool -> RuleM rep Bool
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

-- If a SegRed contains two reduction operations that have the same
-- vector shape, merge them together.  This saves on communication
-- overhead, but can in principle lead to more shared memory usage.
topDownSegOp SymbolTable rep
_ (Pat [PatElem (LetDec rep)]
pes) StmAux (ExpDec rep)
_ (SegRed SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
ts GBody rep KernelResult
kbody [SegBinOp rep]
ops)
  | [SegBinOp rep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SegBinOp rep]
ops Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1,
    [[(SegBinOp rep,
   [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]]
op_groupings <-
      ((SegBinOp rep,
  [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
 -> (SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
 -> Bool)
-> [(SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
-> [[(SegBinOp rep,
      [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (SegBinOp rep,
 [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
-> (SegBinOp rep,
    [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
-> Bool
forall {rep} {b} {rep} {b}.
(SegBinOp rep, b) -> (SegBinOp rep, b) -> Bool
sameShape ([(SegBinOp rep,
   [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
 -> [[(SegBinOp rep,
       [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]])
-> [(SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
-> [[(SegBinOp rep,
      [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]]
forall a b. (a -> b) -> a -> b
$
        [SegBinOp rep]
-> [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
-> [(SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
forall a b. [a] -> [b] -> [(a, b)]
zip [SegBinOp rep]
ops ([[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
 -> [(SegBinOp rep,
      [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])])
-> [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
-> [(SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
forall a b. (a -> b) -> a -> b
$
          [Int]
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
forall a. [Int] -> [a] -> [[a]]
chunks ((SegBinOp rep -> Int) -> [SegBinOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int)
-> (SegBinOp rep -> [SubExp]) -> SegBinOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral) [SegBinOp rep]
ops) ([(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
 -> [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]])
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
forall a b. (a -> b) -> a -> b
$
            [PatElem (LetDec rep)]
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [PatElem (LetDec rep)]
red_pes [TypeBase Shape NoMode]
red_ts [KernelResult]
red_res,
    ([(SegBinOp rep,
   [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
 -> Bool)
-> [[(SegBinOp rep,
      [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]]
-> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (Int -> Bool)
-> ([(SegBinOp rep,
      [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
    -> Int)
-> [(SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [(SegBinOp rep,
  [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
-> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[(SegBinOp rep,
   [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]]
op_groupings = RuleM rep () -> Rule rep
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM rep () -> Rule rep) -> RuleM rep () -> Rule rep
forall a b. (a -> b) -> a -> b
$ do
      let ([SegBinOp rep]
ops', [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
aux) = [(SegBinOp rep,
  [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
-> ([SegBinOp rep],
    [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(SegBinOp rep,
   [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
 -> ([SegBinOp rep],
     [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]))
-> [(SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
-> ([SegBinOp rep],
    [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]])
forall a b. (a -> b) -> a -> b
$ ([(SegBinOp rep,
   [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
 -> Maybe
      (SegBinOp rep,
       [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]))
-> [[(SegBinOp rep,
      [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]]
-> [(SegBinOp rep,
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [(SegBinOp rep,
  [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]
-> Maybe
     (SegBinOp rep,
      [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
forall {rep} {a}.
Buildable rep =>
[(SegBinOp rep, [a])] -> Maybe (SegBinOp rep, [a])
combineOps [[(SegBinOp rep,
   [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])]]
op_groupings
          ([PatElem (LetDec rep)]
red_pes', [TypeBase Shape NoMode]
red_ts', [KernelResult]
red_res') = [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> ([PatElem (LetDec rep)], [TypeBase Shape NoMode],
    [KernelResult])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
 -> ([PatElem (LetDec rep)], [TypeBase Shape NoMode],
     [KernelResult]))
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> ([PatElem (LetDec rep)], [TypeBase Shape NoMode],
    [KernelResult])
forall a b. (a -> b) -> a -> b
$ [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]]
aux
          pes' :: [PatElem (LetDec rep)]
pes' = [PatElem (LetDec rep)]
red_pes' [PatElem (LetDec rep)]
-> [PatElem (LetDec rep)] -> [PatElem (LetDec rep)]
forall a. [a] -> [a] -> [a]
++ [PatElem (LetDec rep)]
map_pes
          ts' :: [TypeBase Shape NoMode]
ts' = [TypeBase Shape NoMode]
red_ts' [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoMode]
map_ts
          kbody' :: GBody rep KernelResult
kbody' = GBody rep KernelResult
kbody {bodyResult = red_res' ++ map_res}
      Pat (LetDec (Rep (RuleM rep)))
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind ([PatElem (LetDec rep)] -> Pat (LetDec rep)
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem (LetDec rep)]
pes') (Exp (Rep (RuleM rep)) -> RuleM rep ())
-> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall a b. (a -> b) -> a -> b
$ Op (Rep (RuleM rep)) -> Exp (Rep (RuleM rep))
forall rep. Op rep -> Exp rep
Op (Op (Rep (RuleM rep)) -> Exp (Rep (RuleM rep)))
-> Op (Rep (RuleM rep)) -> Exp (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$ SegOp (SegOpLevel rep) rep -> Op rep
forall rep. HasSegOp rep => SegOp (SegOpLevel rep) rep -> Op rep
segOp (SegOp (SegOpLevel rep) rep -> Op rep)
-> SegOp (SegOpLevel rep) rep -> Op rep
forall a b. (a -> b) -> a -> b
$ SegOpLevel rep
-> SegSpace
-> [TypeBase Shape NoMode]
-> GBody rep KernelResult
-> [SegBinOp rep]
-> SegOp (SegOpLevel rep) rep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
SegRed SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
ts' GBody rep KernelResult
kbody' [SegBinOp rep]
ops'
  where
    ([PatElem (LetDec rep)]
red_pes, [PatElem (LetDec rep)]
map_pes) = Int
-> [PatElem (LetDec rep)]
-> ([PatElem (LetDec rep)], [PatElem (LetDec rep)])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SegBinOp rep] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp rep]
ops) [PatElem (LetDec rep)]
pes
    ([TypeBase Shape NoMode]
red_ts, [TypeBase Shape NoMode]
map_ts) = Int
-> [TypeBase Shape NoMode]
-> ([TypeBase Shape NoMode], [TypeBase Shape NoMode])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SegBinOp rep] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp rep]
ops) [TypeBase Shape NoMode]
ts
    ([KernelResult]
red_res, [KernelResult]
map_res) = Int -> [KernelResult] -> ([KernelResult], [KernelResult])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SegBinOp rep] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp rep]
ops) ([KernelResult] -> ([KernelResult], [KernelResult]))
-> [KernelResult] -> ([KernelResult], [KernelResult])
forall a b. (a -> b) -> a -> b
$ GBody rep KernelResult -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult GBody rep KernelResult
kbody

    sameShape :: (SegBinOp rep, b) -> (SegBinOp rep, b) -> Bool
sameShape (SegBinOp rep
op1, b
_) (SegBinOp rep
op2, b
_) =
      SegBinOp rep -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp rep
op1 Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
== SegBinOp rep -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp rep
op2
        Bool -> Bool -> Bool
&& Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (SegBinOp rep -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp rep
op1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0

    combineOps :: [(SegBinOp rep, [a])] -> Maybe (SegBinOp rep, [a])
combineOps [] = Maybe (SegBinOp rep, [a])
forall a. Maybe a
Nothing
    combineOps ((SegBinOp rep, [a])
x : [(SegBinOp rep, [a])]
xs) = (SegBinOp rep, [a]) -> Maybe (SegBinOp rep, [a])
forall a. a -> Maybe a
Just ((SegBinOp rep, [a]) -> Maybe (SegBinOp rep, [a]))
-> (SegBinOp rep, [a]) -> Maybe (SegBinOp rep, [a])
forall a b. (a -> b) -> a -> b
$ ((SegBinOp rep, [a]) -> (SegBinOp rep, [a]) -> (SegBinOp rep, [a]))
-> (SegBinOp rep, [a])
-> [(SegBinOp rep, [a])]
-> (SegBinOp rep, [a])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (SegBinOp rep, [a]) -> (SegBinOp rep, [a]) -> (SegBinOp rep, [a])
forall {rep} {a}.
Buildable rep =>
(SegBinOp rep, [a]) -> (SegBinOp rep, [a]) -> (SegBinOp rep, [a])
combine (SegBinOp rep, [a])
x [(SegBinOp rep, [a])]
xs

    combine :: (SegBinOp rep, [a]) -> (SegBinOp rep, [a]) -> (SegBinOp rep, [a])
combine (SegBinOp rep
op1, [a]
op1_aux) (SegBinOp rep
op2, [a]
op2_aux) =
      let lam1 :: Lambda rep
lam1 = SegBinOp rep -> Lambda rep
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp rep
op1
          lam2 :: Lambda rep
lam2 = SegBinOp rep -> Lambda rep
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp rep
op2
          ([Param (LParamInfo rep)]
op1_xparams, [Param (LParamInfo rep)]
op1_yparams) =
            Int
-> [Param (LParamInfo rep)]
-> ([Param (LParamInfo rep)], [Param (LParamInfo rep)])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp rep
op1)) ([Param (LParamInfo rep)]
 -> ([Param (LParamInfo rep)], [Param (LParamInfo rep)]))
-> [Param (LParamInfo rep)]
-> ([Param (LParamInfo rep)], [Param (LParamInfo rep)])
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Param (LParamInfo rep)]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda rep
lam1
          ([Param (LParamInfo rep)]
op2_xparams, [Param (LParamInfo rep)]
op2_yparams) =
            Int
-> [Param (LParamInfo rep)]
-> ([Param (LParamInfo rep)], [Param (LParamInfo rep)])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp rep
op2)) ([Param (LParamInfo rep)]
 -> ([Param (LParamInfo rep)], [Param (LParamInfo rep)]))
-> [Param (LParamInfo rep)]
-> ([Param (LParamInfo rep)], [Param (LParamInfo rep)])
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Param (LParamInfo rep)]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda rep
lam2
          lam :: Lambda rep
lam =
            Lambda
              { lambdaParams :: [Param (LParamInfo rep)]
lambdaParams =
                  [Param (LParamInfo rep)]
op1_xparams
                    [Param (LParamInfo rep)]
-> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. [a] -> [a] -> [a]
++ [Param (LParamInfo rep)]
op2_xparams
                    [Param (LParamInfo rep)]
-> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. [a] -> [a] -> [a]
++ [Param (LParamInfo rep)]
op1_yparams
                    [Param (LParamInfo rep)]
-> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. [a] -> [a] -> [a]
++ [Param (LParamInfo rep)]
op2_yparams,
                lambdaReturnType :: [TypeBase Shape NoMode]
lambdaReturnType = Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam1 [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam2,
                lambdaBody :: Body rep
lambdaBody =
                  Stms rep -> [SubExpRes] -> Body rep
forall res. IsResult res => Stms rep -> [res] -> GBody rep res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody (Body rep -> Stms rep
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda rep -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam1) Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Body rep -> Stms rep
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda rep -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam2)) ([SubExpRes] -> Body rep) -> [SubExpRes] -> Body rep
forall a b. (a -> b) -> a -> b
$
                    Body rep -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda rep -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam1) [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> Body rep -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda rep -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam2)
              }
       in ( SegBinOp
              { segBinOpComm :: Commutativity
segBinOpComm = SegBinOp rep -> Commutativity
forall rep. SegBinOp rep -> Commutativity
segBinOpComm SegBinOp rep
op1 Commutativity -> Commutativity -> Commutativity
forall a. Semigroup a => a -> a -> a
<> SegBinOp rep -> Commutativity
forall rep. SegBinOp rep -> Commutativity
segBinOpComm SegBinOp rep
op2,
                segBinOpLambda :: Lambda rep
segBinOpLambda = Lambda rep
lam,
                segBinOpNeutral :: [SubExp]
segBinOpNeutral = SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp rep
op1 [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp rep
op2,
                segBinOpShape :: Shape
segBinOpShape = SegBinOp rep -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp rep
op1 -- Same as shape of op2 due to the grouping.
              },
            [a]
op1_aux [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
op2_aux
          )

-- Remove unused kernelBody result in SegScan.
topDownSegOp SymbolTable rep
_ Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux SegOp (SegOpLevel rep) rep
op
  | SegScan SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
ts GBody rep KernelResult
kbody [SegBinOp rep]
seg_op SegPostOp rep
post_op <- SegOp (SegOpLevel rep) rep
op,
    -- Figure out which of the names in 'pat' are used...
    Just (GBody rep KernelResult
new_kbody, [TypeBase Shape NoMode]
new_ts, RuleM rep (SegPostOp rep)
m_new_post_op) <-
      GBody rep KernelResult
-> [TypeBase Shape NoMode]
-> [SegBinOp rep]
-> SegPostOp rep
-> Maybe
     (GBody rep KernelResult, [TypeBase Shape NoMode],
      RuleM rep (SegPostOp rep))
forall {rep} {m :: * -> *} {rep}.
(MonadFreshNames m, Buildable rep) =>
KernelBody rep
-> [TypeBase Shape NoMode]
-> [SegBinOp rep]
-> SegPostOp rep
-> Maybe
     (KernelBody rep, [TypeBase Shape NoMode], m (SegPostOp rep))
newKbodyPostOp GBody rep KernelResult
kbody [TypeBase Shape NoMode]
ts [SegBinOp rep]
seg_op SegPostOp rep
post_op = RuleM rep () -> Rule rep
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM rep () -> Rule rep) -> RuleM rep () -> Rule rep
forall a b. (a -> b) -> a -> b
$ do
      new_post_op <- RuleM rep (SegPostOp rep)
m_new_post_op
      auxing aux
        . letBind pat
        . Op
        . segOp
        $ SegScan lvl space new_ts new_kbody seg_op new_post_op
  where
    newKbodyPostOp :: KernelBody rep
-> [TypeBase Shape NoMode]
-> [SegBinOp rep]
-> SegPostOp rep
-> Maybe
     (KernelBody rep, [TypeBase Shape NoMode], m (SegPostOp rep))
newKbodyPostOp KernelBody rep
kbody [TypeBase Shape NoMode]
ts [SegBinOp rep]
seg_op SegPostOp rep
post_op =
      if [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
-> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
sub_map_res_ts_pars
        then Maybe (KernelBody rep, [TypeBase Shape NoMode], m (SegPostOp rep))
forall a. Maybe a
Nothing
        else (KernelBody rep, [TypeBase Shape NoMode], m (SegPostOp rep))
-> Maybe
     (KernelBody rep, [TypeBase Shape NoMode], m (SegPostOp rep))
forall a. a -> Maybe a
Just (KernelBody rep
new_kbody, [TypeBase Shape NoMode]
new_ts, m (SegPostOp rep)
new_post_op)
      where
        res :: [KernelResult]
res = KernelBody rep -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody rep
kbody
        post_lam :: Lambda rep
post_lam = SegPostOp rep -> Lambda rep
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp rep
post_op
        pars :: [Param (LParamInfo rep)]
pars = Lambda rep -> [Param (LParamInfo rep)]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda rep
post_lam

        mkBind :: TypeBase Shape NoMode -> Param dec -> KernelResult -> Stm rep
mkBind TypeBase Shape NoMode
t Param dec
p KernelResult
r =
          [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet
            [VName -> TypeBase Shape NoMode -> Ident
Ident (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
p) TypeBase Shape NoMode
t]
            (BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ KernelResult -> SubExp
kernelResultSubExp KernelResult
r)

        new_kbody :: KernelBody rep
new_kbody = KernelBody rep
kbody {bodyResult = new_res}
        new_post_op :: m (SegPostOp rep)
new_post_op = do
          let sub_res :: [KernelResult]
sub_res = ((KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)
 -> KernelResult)
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> [KernelResult]
forall a b. (a -> b) -> [a] -> [b]
map (\(KernelResult
r, TypeBase Shape NoMode
_, Param (LParamInfo rep)
_, Names
_) -> KernelResult
r) [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
sub_map_res_ts_pars
          temp_body <- KernelBody rep -> m (KernelBody rep)
forall rep res (m :: * -> *).
(Renameable rep, Rename res, MonadFreshNames m) =>
GBody rep res -> m (GBody rep res)
renameBody (KernelBody rep -> m (KernelBody rep))
-> KernelBody rep -> m (KernelBody rep)
forall a b. (a -> b) -> a -> b
$ KernelBody rep
kbody {bodyResult = sub_res}
          let new_binds =
                [Stm rep] -> Stms rep
forall rep. [Stm rep] -> Stms rep
stmsFromList
                  ([Stm rep] -> Stms rep) -> [Stm rep] -> Stms rep
forall a b. (a -> b) -> a -> b
$ ((KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)
 -> KernelResult -> Stm rep)
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> [KernelResult]
-> [Stm rep]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
                    (\(KernelResult
_, TypeBase Shape NoMode
t, Param (LParamInfo rep)
p, Names
_) KernelResult
r -> TypeBase Shape NoMode
-> Param (LParamInfo rep) -> KernelResult -> Stm rep
forall {rep} {dec}.
Buildable rep =>
TypeBase Shape NoMode -> Param dec -> KernelResult -> Stm rep
mkBind TypeBase Shape NoMode
t Param (LParamInfo rep)
p KernelResult
r)
                    [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
sub_map_res_ts_pars
                  ([KernelResult] -> [Stm rep]) -> [KernelResult] -> [Stm rep]
forall a b. (a -> b) -> a -> b
$ KernelBody rep -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody rep
temp_body
          pure $
            SegPostOp $
              Lambda
                { lambdaParams = new_pars,
                  lambdaBody =
                    mkBody
                      (bodyStms temp_body <> new_binds <> bodyStms (lambdaBody post_lam))
                      (bodyResult (lambdaBody post_lam)),
                  lambdaReturnType = lambdaReturnType post_lam
                }

        scan_deps :: Names
scan_deps = [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> [Names] -> Names
forall a b. (a -> b) -> a -> b
$ (\(KernelResult
_, TypeBase Shape NoMode
_, Param (LParamInfo rep)
_, Names
d) -> Names
d) ((KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)
 -> Names)
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> [Names]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
scan_res_ts_pars
        deps :: [Names]
deps = AliasTable -> KernelBody rep -> [Names]
forall rep. ASTRep rep => AliasTable -> KernelBody rep -> [Names]
kernelBodyDependencies AliasTable
forall a. Monoid a => a
mempty KernelBody rep
kbody

        isNotReturns :: KernelResult -> Bool
isNotReturns (Returns {}) = Bool
False
        isNotReturns KernelResult
_ = Bool
True

        ([KernelResult]
new_res, [TypeBase Shape NoMode]
new_ts, [Param (LParamInfo rep)]
new_pars, [Names]
_) =
          [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
-> ([KernelResult], [TypeBase Shape NoMode],
    [Param (LParamInfo rep)], [Names])
forall a b c d. [(a, b, c, d)] -> ([a], [b], [c], [d])
L.unzip4 ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
   Names)]
 -> ([KernelResult], [TypeBase Shape NoMode],
     [Param (LParamInfo rep)], [Names]))
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> ([KernelResult], [TypeBase Shape NoMode],
    [Param (LParamInfo rep)], [Names])
forall a b. (a -> b) -> a -> b
$ [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
scan_res_ts_pars [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
forall a. Semigroup a => a -> a -> a
<> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
new_map_res_ts_pars
        ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
new_map_res_ts_pars, [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
sub_map_res_ts_pars) =
          ((KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)
 -> Bool)
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
      Names)],
    [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
      Names)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition
            ( \(KernelResult
r, TypeBase Shape NoMode
t, Param (LParamInfo rep)
_, Names
d) ->
                KernelResult -> Bool
isNotReturns KernelResult
r Bool -> Bool -> Bool
|| Names
d Names -> Names -> Bool
`namesIntersect` Names
scan_deps Bool -> Bool -> Bool
|| TypeBase Shape NoMode -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc TypeBase Shape NoMode
t
            )
            [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
map_res_ts_pars
        ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
scan_res_ts_pars, [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
  Names)]
map_res_ts_pars) =
          Int
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
      Names)],
    [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
      Names)])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SegBinOp rep] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp rep]
seg_op) ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
   Names)]
 -> ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
       Names)],
     [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
       Names)]))
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
-> ([(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
      Names)],
    [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
      Names)])
forall a b. (a -> b) -> a -> b
$ [KernelResult]
-> [TypeBase Shape NoMode]
-> [Param (LParamInfo rep)]
-> [Names]
-> [(KernelResult, TypeBase Shape NoMode, Param (LParamInfo rep),
     Names)]
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
L.zip4 [KernelResult]
res [TypeBase Shape NoMode]
ts [Param (LParamInfo rep)]
pars [Names]
deps
topDownSegOp SymbolTable rep
_ Pat (LetDec rep)
_ StmAux (ExpDec rep)
_ SegOp (SegOpLevel rep) rep
_ = Rule rep
forall rep. Rule rep
Skip

-- A convenient way of operating on the type and body of a SegOp,
-- without worrying about exactly what kind it is.
segOpGuts ::
  SegOp (SegOpLevel rep) rep ->
  ( [Type],
    KernelBody rep,
    Int,
    [Type] -> KernelBody rep -> SegOp (SegOpLevel rep) rep
  )
segOpGuts :: forall rep.
SegOp (SegOpLevel rep) rep
-> ([TypeBase Shape NoMode], KernelBody rep, Int,
    [TypeBase Shape NoMode]
    -> KernelBody rep -> SegOp (SegOpLevel rep) rep)
segOpGuts (SegMap SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
kts KernelBody rep
body) =
  ([TypeBase Shape NoMode]
kts, KernelBody rep
body, Int
0, SegOpLevel rep
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> SegOp (SegOpLevel rep) rep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> SegOp lvl rep
SegMap SegOpLevel rep
lvl SegSpace
space)
segOpGuts (SegScan SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
kts KernelBody rep
body [SegBinOp rep]
ops SegPostOp rep
post_op) =
  ([TypeBase Shape NoMode]
kts, KernelBody rep
body, Int
0, \[TypeBase Shape NoMode]
t KernelBody rep
b -> SegOpLevel rep
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegPostOp rep
-> SegOp (SegOpLevel rep) rep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegPostOp rep
-> SegOp lvl rep
SegScan SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
t KernelBody rep
b [SegBinOp rep]
ops SegPostOp rep
post_op)
segOpGuts (SegRed SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
kts KernelBody rep
body [SegBinOp rep]
ops) =
  ([TypeBase Shape NoMode]
kts, KernelBody rep
body, [SegBinOp rep] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp rep]
ops, \[TypeBase Shape NoMode]
t KernelBody rep
b -> SegOpLevel rep
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp (SegOpLevel rep) rep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
SegRed SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
t KernelBody rep
b [SegBinOp rep]
ops)
segOpGuts (SegHist SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
kts KernelBody rep
body [HistOp rep]
ops) =
  ([TypeBase Shape NoMode]
kts, KernelBody rep
body, [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (HistOp rep -> Int) -> [HistOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([VName] -> Int) -> (HistOp rep -> [VName]) -> HistOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. HistOp rep -> [VName]
forall rep. HistOp rep -> [VName]
histDest) [HistOp rep]
ops, \[TypeBase Shape NoMode]
t KernelBody rep
b -> SegOpLevel rep
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [HistOp rep]
-> SegOp (SegOpLevel rep) rep
forall lvl rep.
lvl
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [HistOp rep]
-> SegOp lvl rep
SegHist SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
t KernelBody rep
b [HistOp rep]
ops)

bottomUpSegOp ::
  (Aliased rep, HasSegOp rep, BuilderOps rep) =>
  (ST.SymbolTable rep, UT.UsageTable) ->
  Pat (LetDec rep) ->
  StmAux (ExpDec rep) ->
  SegOp (SegOpLevel rep) rep ->
  Rule rep
-- Because of the postop we have to handle SegScan specially.
bottomUpSegOp :: forall rep.
(Aliased rep, HasSegOp rep, BuilderOps rep) =>
(SymbolTable rep, UsageTable)
-> Pat (LetDec rep)
-> StmAux (ExpDec rep)
-> SegOp (SegOpLevel rep) rep
-> Rule rep
bottomUpSegOp (SymbolTable rep
_vtable, UsageTable
used) (Pat [PatElem (LetDec rep)]
kpes) StmAux (ExpDec rep)
aux (SegScan SegOpLevel rep
lvl SegSpace
space [TypeBase Shape NoMode]
kts KernelBody rep
body [SegBinOp rep]
ops (SegPostOp Lambda rep
postlam))
  | -- Remove dead results from the postlam.
    ([PatElem (LetDec rep)]
kpes', [SubExpRes]
res') <-
      [(PatElem (LetDec rep), SubExpRes)]
-> ([PatElem (LetDec rep)], [SubExpRes])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PatElem (LetDec rep), SubExpRes)]
 -> ([PatElem (LetDec rep)], [SubExpRes]))
-> [(PatElem (LetDec rep), SubExpRes)]
-> ([PatElem (LetDec rep)], [SubExpRes])
forall a b. (a -> b) -> a -> b
$ ((PatElem (LetDec rep), SubExpRes) -> Bool)
-> [(PatElem (LetDec rep), SubExpRes)]
-> [(PatElem (LetDec rep), SubExpRes)]
forall a. (a -> Bool) -> [a] -> [a]
filter (PatElem (LetDec rep), SubExpRes) -> Bool
keep ([(PatElem (LetDec rep), SubExpRes)]
 -> [(PatElem (LetDec rep), SubExpRes)])
-> [(PatElem (LetDec rep), SubExpRes)]
-> [(PatElem (LetDec rep), SubExpRes)]
forall a b. (a -> b) -> a -> b
$ [PatElem (LetDec rep)]
-> [SubExpRes] -> [(PatElem (LetDec rep), SubExpRes)]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatElem (LetDec rep)]
kpes (GBody rep SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
postlam)),
    [PatElem (LetDec rep)]
kpes [PatElem (LetDec rep)] -> [PatElem (LetDec rep)] -> Bool
forall a. Eq a => a -> a -> Bool
/= [PatElem (LetDec rep)]
kpes' = RuleM rep () -> Rule rep
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM rep () -> Rule rep) -> RuleM rep () -> Rule rep
forall a b. (a -> b) -> a -> b
$ do
      postlam' <- [LParam (Rep (RuleM rep))]
-> RuleM rep [SubExpRes] -> RuleM rep (Lambda (Rep (RuleM rep)))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m [SubExpRes] -> m (Lambda (Rep m))
mkLambda (Lambda rep -> [LParam rep]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda rep
postlam) (RuleM rep [SubExpRes] -> RuleM rep (Lambda (Rep (RuleM rep))))
-> RuleM rep [SubExpRes] -> RuleM rep (Lambda (Rep (RuleM rep)))
forall a b. (a -> b) -> a -> b
$ do
        Stms (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms (Stms (Rep (RuleM rep)) -> RuleM rep ())
-> Stms (Rep (RuleM rep)) -> RuleM rep ()
forall a b. (a -> b) -> a -> b
$ GBody (Rep (RuleM rep)) SubExpRes -> Stms (Rep (RuleM rep))
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody (Rep (RuleM rep)) SubExpRes -> Stms (Rep (RuleM rep)))
-> GBody (Rep (RuleM rep)) SubExpRes -> Stms (Rep (RuleM rep))
forall a b. (a -> b) -> a -> b
$ Lambda (Rep (RuleM rep)) -> GBody (Rep (RuleM rep)) SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
Lambda (Rep (RuleM rep))
postlam
        [SubExpRes] -> RuleM rep [SubExpRes]
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExpRes]
res'
      addStm $ Let (Pat kpes') aux $ Op $ segOp $ SegScan lvl space kts body ops (SegPostOp postlam')
  | Bool
otherwise = Rule rep
forall rep. Rule rep
Skip
  where
    keep :: (PatElem (LetDec rep), SubExpRes) -> Bool
keep (PatElem (LetDec rep)
pe, SubExpRes
_) = PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe VName -> UsageTable -> Bool
`UT.used` UsageTable
used
bottomUpSegOp (SymbolTable rep
_vtable, UsageTable
used) (Pat [PatElem (LetDec rep)]
kpes) StmAux (ExpDec rep)
dec SegOp (SegOpLevel rep) rep
segop
  -- Remove dead results. This is a bit tricky to do with reduction results, so
  -- we only deal with map results for now.
  | ([Int]
_, [PatElem (LetDec rep)]
kpes', [TypeBase Shape NoMode]
kts', [KernelResult]
kres') <- [(Int, PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> ([Int], [PatElem (LetDec rep)], [TypeBase Shape NoMode],
    [KernelResult])
forall a b c d. [(a, b, c, d)] -> ([a], [b], [c], [d])
unzip4 ([(Int, PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
 -> ([Int], [PatElem (LetDec rep)], [TypeBase Shape NoMode],
     [KernelResult]))
-> [(Int, PatElem (LetDec rep), TypeBase Shape NoMode,
     KernelResult)]
-> ([Int], [PatElem (LetDec rep)], [TypeBase Shape NoMode],
    [KernelResult])
forall a b. (a -> b) -> a -> b
$ ((Int, PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)
 -> Bool)
-> [(Int, PatElem (LetDec rep), TypeBase Shape NoMode,
     KernelResult)]
-> [(Int, PatElem (LetDec rep), TypeBase Shape NoMode,
     KernelResult)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int, PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)
-> Bool
keep ([(Int, PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
 -> [(Int, PatElem (LetDec rep), TypeBase Shape NoMode,
      KernelResult)])
-> [(Int, PatElem (LetDec rep), TypeBase Shape NoMode,
     KernelResult)]
-> [(Int, PatElem (LetDec rep), TypeBase Shape NoMode,
     KernelResult)]
forall a b. (a -> b) -> a -> b
$ [Int]
-> [PatElem (LetDec rep)]
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> [(Int, PatElem (LetDec rep), TypeBase Shape NoMode,
     KernelResult)]
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 [Int
0 ..] [PatElem (LetDec rep)]
kpes [TypeBase Shape NoMode]
kts [KernelResult]
kres,
    [PatElem (LetDec rep)]
kpes' [PatElem (LetDec rep)] -> [PatElem (LetDec rep)] -> Bool
forall a. Eq a => a -> a -> Bool
/= [PatElem (LetDec rep)]
kpes = RuleM rep () -> Rule rep
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM rep () -> Rule rep) -> RuleM rep () -> Rule rep
forall a b. (a -> b) -> a -> b
$ do
      kbody' <- Scope rep
-> RuleM rep (KernelBody rep) -> RuleM rep (KernelBody rep)
forall a. Scope rep -> RuleM rep a -> RuleM rep a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (SegSpace -> Scope rep
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (RuleM rep (KernelBody rep) -> RuleM rep (KernelBody rep))
-> RuleM rep (KernelBody rep) -> RuleM rep (KernelBody rep)
forall a b. (a -> b) -> a -> b
$ Stms (Rep (RuleM rep))
-> [KernelResult]
-> RuleM rep (GBody (Rep (RuleM rep)) KernelResult)
forall res.
IsResult res =>
Stms (Rep (RuleM rep))
-> [res] -> RuleM rep (GBody (Rep (RuleM rep)) res)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
Stms (Rep m) -> [res] -> m (GBody (Rep m) res)
mkBodyM Stms rep
Stms (Rep (RuleM rep))
kstms [KernelResult]
kres'
      addStm $ Let (Pat kpes') dec $ Op $ segOp $ mk_segop kts' kbody'
  where
    space :: SegSpace
space = SegOp (SegOpLevel rep) rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp (SegOpLevel rep) rep
segop
    ([TypeBase Shape NoMode]
kts, Body BodyDec rep
_ Stms rep
kstms [KernelResult]
kres, Int
num_nonmap_results, [TypeBase Shape NoMode]
-> KernelBody rep -> SegOp (SegOpLevel rep) rep
mk_segop) =
      SegOp (SegOpLevel rep) rep
-> ([TypeBase Shape NoMode], KernelBody rep, Int,
    [TypeBase Shape NoMode]
    -> KernelBody rep -> SegOp (SegOpLevel rep) rep)
forall rep.
SegOp (SegOpLevel rep) rep
-> ([TypeBase Shape NoMode], KernelBody rep, Int,
    [TypeBase Shape NoMode]
    -> KernelBody rep -> SegOp (SegOpLevel rep) rep)
segOpGuts SegOp (SegOpLevel rep) rep
segop

    keep :: (Int, PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)
-> Bool
keep (Int
i, PatElem (LetDec rep)
pe, TypeBase Shape NoMode
_, KernelResult
_) =
      Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
num_nonmap_results Bool -> Bool -> Bool
|| PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe VName -> UsageTable -> Bool
`UT.used` UsageTable
used
bottomUpSegOp (SymbolTable rep
vtable, UsageTable
_used) (Pat [PatElem (LetDec rep)]
kpes) StmAux (ExpDec rep)
dec SegOp (SegOpLevel rep) rep
segop = RuleM rep () -> Rule rep
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM rep () -> Rule rep) -> RuleM rep () -> Rule rep
forall a b. (a -> b) -> a -> b
$ do
  -- Iterate through the bindings.  For each, we check whether it is
  -- in kres and can be moved outside.  If so, we remove it from kres
  -- and kpes and make it a binding outside.  We have to be careful
  -- not to remove anything that is passed on to a scan/map/histogram
  -- operation.  Fortunately, these are always first in the result
  -- list.
  (kpes', kts', kres', kstms') <-
    Scope rep
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
forall a. Scope rep -> RuleM rep a -> RuleM rep a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (SegSpace -> Scope rep
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (RuleM
   rep
   ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
    Stms rep)
 -> RuleM
      rep
      ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
       Stms rep))
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
forall a b. (a -> b) -> a -> b
$
      (([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
  Stms rep)
 -> Stm rep
 -> RuleM
      rep
      ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
       Stms rep))
-> ([PatElem (LetDec rep)], [TypeBase Shape NoMode],
    [KernelResult], Stms rep)
-> Stms rep
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
 Stms rep)
-> Stm rep
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
distribute ([PatElem (LetDec rep)]
kpes, [TypeBase Shape NoMode]
kts, [KernelResult]
kres, Stms rep
forall a. Monoid a => a
mempty) Stms rep
kstms

  when (kpes' == kpes) cannotSimplify

  kbody' <-
    localScope (scopeOfSegSpace space) $ mkBodyM kstms' kres'

  addStm $ Let (Pat kpes') dec $ Op $ segOp $ mk_segop kts' kbody'
  where
    ([TypeBase Shape NoMode]
kts, Body BodyDec rep
_ Stms rep
kstms [KernelResult]
kres, Int
num_nonmap_results, [TypeBase Shape NoMode]
-> KernelBody rep -> SegOp (SegOpLevel rep) rep
mk_segop) =
      SegOp (SegOpLevel rep) rep
-> ([TypeBase Shape NoMode], KernelBody rep, Int,
    [TypeBase Shape NoMode]
    -> KernelBody rep -> SegOp (SegOpLevel rep) rep)
forall rep.
SegOp (SegOpLevel rep) rep
-> ([TypeBase Shape NoMode], KernelBody rep, Int,
    [TypeBase Shape NoMode]
    -> KernelBody rep -> SegOp (SegOpLevel rep) rep)
segOpGuts SegOp (SegOpLevel rep) rep
segop
    free_in_kstms :: Names
free_in_kstms = (Stm rep -> Names) -> Stms rep -> Names
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm rep -> Names
forall a. FreeIn a => a -> Names
freeIn Stms rep
kstms
    space :: SegSpace
space = SegOp (SegOpLevel rep) rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp (SegOpLevel rep) rep
segop

    sliceWithGtidsFixed :: Stm rep -> Maybe (Slice SubExp, VName)
sliceWithGtidsFixed Stm rep
stm
      | Let Pat (LetDec rep)
_ StmAux (ExpDec rep)
aux (BasicOp (Index VName
arr Slice SubExp
slice)) <- Stm rep
stm,
        [DimIndex SubExp]
space_slice <- ((VName, SubExp) -> DimIndex SubExp)
-> [(VName, SubExp)] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp)
-> ((VName, SubExp) -> SubExp)
-> (VName, SubExp)
-> DimIndex SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VName -> SubExp
Var (VName -> SubExp)
-> ((VName, SubExp) -> VName) -> (VName, SubExp) -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst) ([(VName, SubExp)] -> [DimIndex SubExp])
-> [(VName, SubExp)] -> [DimIndex SubExp]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space,
        [DimIndex SubExp]
space_slice [DimIndex SubExp] -> [DimIndex SubExp] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice,
        Slice SubExp
remaining_slice <- [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ Int -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. Int -> [a] -> [a]
drop ([DimIndex SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex SubExp]
space_slice) (Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice),
        (VName -> Bool) -> Names -> Bool
allNames (Maybe (Entry rep) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Entry rep) -> Bool)
-> (VName -> Maybe (Entry rep)) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (VName -> SymbolTable rep -> Maybe (Entry rep))
-> SymbolTable rep -> VName -> Maybe (Entry rep)
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> SymbolTable rep -> Maybe (Entry rep)
forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
ST.lookup SymbolTable rep
vtable) (Names -> Bool) -> Names -> Bool
forall a b. (a -> b) -> a -> b
$
          VName -> Names
forall a. FreeIn a => a -> Names
freeIn VName
arr Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Slice SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn Slice SubExp
remaining_slice Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Certs -> Names
forall a. FreeIn a => a -> Names
freeIn (StmAux (ExpDec rep) -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec rep)
aux) =
          (Slice SubExp, VName) -> Maybe (Slice SubExp, VName)
forall a. a -> Maybe a
Just (Slice SubExp
remaining_slice, VName
arr)
      | Bool
otherwise =
          Maybe (Slice SubExp, VName)
forall a. Maybe a
Nothing

    distribute :: ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
 Stms rep)
-> Stm rep
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
distribute ([PatElem (LetDec rep)]
kpes', [TypeBase Shape NoMode]
kts', [KernelResult]
kres', Stms rep
kstms') Stm rep
stm
      | Let (Pat [PatElem (LetDec rep)
pe]) StmAux (ExpDec rep)
_ Exp rep
_ <- Stm rep
stm,
        Just (Slice [DimIndex SubExp]
remaining_slice, VName
arr) <- Stm rep -> Maybe (Slice SubExp, VName)
sliceWithGtidsFixed Stm rep
stm,
        Just (PatElem (LetDec rep)
kpe, [PatElem (LetDec rep)]
kpes'', [TypeBase Shape NoMode]
kts'', [KernelResult]
kres'') <- [PatElem (LetDec rep)]
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> PatElem (LetDec rep)
-> Maybe
     (PatElem (LetDec rep), [PatElem (LetDec rep)],
      [TypeBase Shape NoMode], [KernelResult])
isResult [PatElem (LetDec rep)]
kpes' [TypeBase Shape NoMode]
kts' [KernelResult]
kres' PatElem (LetDec rep)
pe = do
          let outer_slice :: [DimIndex SubExp]
outer_slice =
                (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map
                  ( \SubExp
d ->
                      SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)) SubExp
d (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64))
                  )
                  ([SubExp] -> [DimIndex SubExp]) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [SubExp]
segSpaceDims SegSpace
space
              index :: PatElem (LetDec rep) -> RuleM rep ()
index PatElem (LetDec rep)
kpe' =
                [VName] -> Exp (Rep (RuleM rep)) -> RuleM rep ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
kpe'] (Exp rep -> RuleM rep ())
-> (Slice SubExp -> Exp rep) -> Slice SubExp -> RuleM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep)
-> (Slice SubExp -> BasicOp) -> Slice SubExp -> Exp rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> RuleM rep ()) -> Slice SubExp -> RuleM rep ()
forall a b. (a -> b) -> a -> b
$
                  [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$
                    [DimIndex SubExp]
outer_slice [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. Semigroup a => a -> a -> a
<> [DimIndex SubExp]
remaining_slice
          precopy <- Name -> RuleM rep VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> RuleM rep VName) -> Name -> RuleM rep VName
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
kpe) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_precopy"
          index kpe {patElemName = precopy}
          letBindNames [patElemName kpe] $ BasicOp $ Replicate mempty $ Var precopy
          pure
            ( kpes'',
              kts'',
              kres'',
              if patElemName pe `nameIn` free_in_kstms
                then kstms' <> oneStm stm
                else kstms'
            )
    distribute ([PatElem (LetDec rep)]
kpes', [TypeBase Shape NoMode]
kts', [KernelResult]
kres', Stms rep
kstms') Stm rep
stm =
      ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
 Stms rep)
-> RuleM
     rep
     ([PatElem (LetDec rep)], [TypeBase Shape NoMode], [KernelResult],
      Stms rep)
forall a. a -> RuleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PatElem (LetDec rep)]
kpes', [TypeBase Shape NoMode]
kts', [KernelResult]
kres', Stms rep
kstms' Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Stm rep -> Stms rep
forall rep. Stm rep -> Stms rep
oneStm Stm rep
stm)

    isResult :: [PatElem (LetDec rep)]
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> PatElem (LetDec rep)
-> Maybe
     (PatElem (LetDec rep), [PatElem (LetDec rep)],
      [TypeBase Shape NoMode], [KernelResult])
isResult [PatElem (LetDec rep)]
kpes' [TypeBase Shape NoMode]
kts' [KernelResult]
kres' PatElem (LetDec rep)
pe =
      case ((PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)
 -> Bool)
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> ([(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)],
    [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult) -> Bool
matches ([(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
 -> ([(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)],
     [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]))
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> ([(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)],
    [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
forall a b. (a -> b) -> a -> b
$ [PatElem (LetDec rep)]
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [PatElem (LetDec rep)]
kpes' [TypeBase Shape NoMode]
kts' [KernelResult]
kres' of
        ([(PatElem (LetDec rep)
kpe, TypeBase Shape NoMode
_, KernelResult
_)], [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
kpes_and_kres)
          | Just Int
i <- PatElem (LetDec rep) -> [PatElem (LetDec rep)] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex PatElem (LetDec rep)
kpe [PatElem (LetDec rep)]
kpes,
            Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
num_nonmap_results,
            ([PatElem (LetDec rep)]
kpes'', [TypeBase Shape NoMode]
kts'', [KernelResult]
kres'') <- [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
-> ([PatElem (LetDec rep)], [TypeBase Shape NoMode],
    [KernelResult])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)]
kpes_and_kres ->
              (PatElem (LetDec rep), [PatElem (LetDec rep)],
 [TypeBase Shape NoMode], [KernelResult])
-> Maybe
     (PatElem (LetDec rep), [PatElem (LetDec rep)],
      [TypeBase Shape NoMode], [KernelResult])
forall a. a -> Maybe a
Just (PatElem (LetDec rep)
kpe, [PatElem (LetDec rep)]
kpes'', [TypeBase Shape NoMode]
kts'', [KernelResult]
kres'')
        ([(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)],
 [(PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)])
_ -> Maybe
  (PatElem (LetDec rep), [PatElem (LetDec rep)],
   [TypeBase Shape NoMode], [KernelResult])
forall a. Maybe a
Nothing
      where
        matches :: (PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult) -> Bool
matches (PatElem (LetDec rep)
_, TypeBase Shape NoMode
_, Returns ResultManifest
_ Certs
_ (Var VName
v)) = VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe
        matches (PatElem (LetDec rep), TypeBase Shape NoMode, KernelResult)
_ = Bool
False

--- Memory

-- | Like 'segOpType', but for memory representations.
segOpReturns ::
  (Mem rep inner, Monad m, HasScope rep m) =>
  SegOp lvl rep ->
  m [ExpReturns]
segOpReturns :: forall rep (inner :: * -> *) (m :: * -> *) lvl.
(Mem rep inner, Monad m, HasScope rep m) =>
SegOp lvl rep -> m [ExpReturns]
segOpReturns k :: SegOp lvl rep
k@(SegMap {}) =
  [ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns]) -> m [ExtType] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOp lvl rep -> m [ExtType]
forall rep (m :: * -> *).
HasScope rep m =>
SegOp lvl rep -> m [ExtType]
forall (op :: * -> *) rep (m :: * -> *).
(TypedOp op, HasScope rep m) =>
op rep -> m [ExtType]
opType SegOp lvl rep
k
segOpReturns k :: SegOp lvl rep
k@(SegRed {}) =
  [ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns]) -> m [ExtType] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOp lvl rep -> m [ExtType]
forall rep (m :: * -> *).
HasScope rep m =>
SegOp lvl rep -> m [ExtType]
forall (op :: * -> *) rep (m :: * -> *).
(TypedOp op, HasScope rep m) =>
op rep -> m [ExtType]
opType SegOp lvl rep
k
segOpReturns k :: SegOp lvl rep
k@(SegScan {}) =
  [ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns]) -> m [ExtType] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOp lvl rep -> m [ExtType]
forall rep (m :: * -> *).
HasScope rep m =>
SegOp lvl rep -> m [ExtType]
forall (op :: * -> *) rep (m :: * -> *).
(TypedOp op, HasScope rep m) =>
op rep -> m [ExtType]
opType SegOp lvl rep
k
segOpReturns (SegHist lvl
_ SegSpace
_ [TypeBase Shape NoMode]
_ KernelBody rep
_ [HistOp rep]
ops) =
  [[ExpReturns]] -> [ExpReturns]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[ExpReturns]] -> [ExpReturns])
-> m [[ExpReturns]] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (HistOp rep -> m [ExpReturns]) -> [HistOp rep] -> m [[ExpReturns]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((VName -> m ExpReturns) -> [VName] -> m [ExpReturns]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> m ExpReturns
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m ExpReturns
varReturns ([VName] -> m [ExpReturns])
-> (HistOp rep -> [VName]) -> HistOp rep -> m [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. HistOp rep -> [VName]
forall rep. HistOp rep -> [VName]
histDest) [HistOp rep]
ops

-- Note [Pruning the index table]
--
-- The 'ST.IndexOp' instance for 'SegMap' answers the question "what is the
-- value of element @is@ of the array produced by this kernel?". It does so by
-- symbolically evaluating the kernel body into a table mapping each variable to
-- a 'ST.Indexed', then looking up the variable returned by the kernel.
--
-- Only one entry of that table is ever read, and an entry depends solely on the
-- entries of the variables that are free in the statement that binds it. Hence
-- the answer is unchanged if we first drop every statement that the result does
-- not transitively depend on, which is what 'dependedOn' does.
--
-- This matters because building an entry is expensive (it constructs a
-- 'PrimExp' for the entire statement), while deciding whether a statement is
-- depended upon is a set membership test and a set union. Kernel bodies are
-- often large while the slice that computes any single result is small, and the
-- simplifier asks this question once per indexing of the array, in every one of
-- its many passes; without pruning, that is quadratic in the size of the kernel.