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

-- | Definition of /Second-Order Array Combinators/ (SOACs), which are
-- the main form of parallelism in the early stages of the compiler.
module Futhark.IR.SOACS.SOAC
  ( SOAC (..),
    ScremaForm (..),
    HistOp (..),
    Scan (..),
    scanResults,
    singleScan,
    Reduce (..),
    redResults,
    singleReduce,

    -- * Utility
    composeBinds,
    scremaType,
    soacType,
    flatMapNonuniform,
    flatMapRowTypes,
    flatMapUniformTypes,
    flatMapSplitMeta,
    flatMapSplitValues,
    typeCheckSOAC,
    mkIdentityLambda,
    nilFn,
    maposcanomapSOAC,
    scanomapSOAC,
    redomapSOAC,
    scanSOAC,
    reduceSOAC,
    mapSOAC,
    isMaposcanomapSOAC,
    isScanomapSOAC,
    isRedomapSOAC,
    isScanSOAC,
    isReduceSOAC,
    isMapSOAC,
    ppScrema,
    ppHist,
    ppStream,

    -- * Generic traversal
    SOACMapper (..),
    identitySOACMapper,
    mapSOACM,
    traverseSOACStms,
  )
where

import Control.Category
import Control.Monad
import Control.Monad.Identity
import Control.Monad.State.Strict
import Control.Monad.Writer
import Data.Either (partitionEithers)
import Data.List (intersperse)
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.Construct
import Futhark.IR
import Futhark.IR.Aliases (Aliases, CanBeAliased (..))
import Futhark.IR.Prop.Aliases
import Futhark.IR.TypeCheck qualified as TC
import Futhark.Optimise.Simplify.Rep
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util (chunks, maybeNth, splitAt3)
import Futhark.Util.Pretty (Doc, align, comma, commasep, docText, parens, ppTuple', pretty, (<+>), (</>))
import Futhark.Util.Pretty qualified as PP
import Prelude hiding (id, (.))

-- | A second-order array combinator (SOAC).
data SOAC rep
  = Stream SubExp [VName] [SubExp] (Lambda rep)
  | -- | @Hist <length> <input arrays> <dest-arrays-and-ops> <bucket fun>@
    --
    -- The final lambda produces indexes and values for the 'HistOp's.
    Hist SubExp [VName] [HistOp rep] (Lambda rep)
  | -- FIXME: this should not be here
    JVP Shape [SubExp] [SubExp] (Lambda rep)
  | -- FIXME: this should not be here
    VJP Shape [SubExp] [SubExp] (Lambda rep)
  | -- FIXME: this should not be here
    WithVJP [SubExp] (Lambda rep) (Lambda rep)
  | -- | A combination of scan, reduction, and map.  The first
    -- t'SubExp' is the size of the input arrays.
    Screma SubExp [VName] (ScremaForm rep)
  | -- | Nonuniform map where the nonuniform results are implicitly concatenated.
    -- The 'ExtLambda' first returns a size @k@, then its value results, each of
    -- which is either /nonuniform/ (an array with @k@ as its outermost size) or
    -- /uniform/ (not mentioning @k@ at all). The size @k@ may be used *only* as
    -- an outermost dimension.
    --
    -- For input of length @n@ returns the following:
    --
    -- * An integer @m@ denoting the size of the data arrays.
    --
    -- * The "shape array" of type @[n]i64@, giving the size of each segment.
    --   This array sums to @m@.
    --
    -- * The "flag array" of type @[m]bool@, indicating for each element when a
    --   new segment begins.
    --
    -- * The "offset array" of type @[n]i64@, indicating for each segment where
    --   its values begin in the data array.
    --
    -- * Finally one array per value result of the lambda, in the same order: a
    --   nonuniform result becomes its concatenation across all iterations, of
    --   outer size @[m]@, and a uniform result the collection of its values, of
    --   outer size @[n]@.
    FlatMap SubExp [VName] (ExtLambda rep)
  deriving (SOAC rep -> SOAC rep -> Bool
(SOAC rep -> SOAC rep -> Bool)
-> (SOAC rep -> SOAC rep -> Bool) -> Eq (SOAC rep)
forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
== :: SOAC rep -> SOAC rep -> Bool
$c/= :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
/= :: SOAC rep -> SOAC rep -> Bool
Eq, Eq (SOAC rep)
Eq (SOAC rep) =>
(SOAC rep -> SOAC rep -> Ordering)
-> (SOAC rep -> SOAC rep -> Bool)
-> (SOAC rep -> SOAC rep -> Bool)
-> (SOAC rep -> SOAC rep -> Bool)
-> (SOAC rep -> SOAC rep -> Bool)
-> (SOAC rep -> SOAC rep -> SOAC rep)
-> (SOAC rep -> SOAC rep -> SOAC rep)
-> Ord (SOAC rep)
SOAC rep -> SOAC rep -> Bool
SOAC rep -> SOAC rep -> Ordering
SOAC rep -> SOAC rep -> SOAC 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 (SOAC rep)
forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Ordering
forall rep. RepTypes rep => SOAC rep -> SOAC rep -> SOAC rep
$ccompare :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Ordering
compare :: SOAC rep -> SOAC rep -> Ordering
$c< :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
< :: SOAC rep -> SOAC rep -> Bool
$c<= :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
<= :: SOAC rep -> SOAC rep -> Bool
$c> :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
> :: SOAC rep -> SOAC rep -> Bool
$c>= :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
>= :: SOAC rep -> SOAC rep -> Bool
$cmax :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> SOAC rep
max :: SOAC rep -> SOAC rep -> SOAC rep
$cmin :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> SOAC rep
min :: SOAC rep -> SOAC rep -> SOAC rep
Ord, Int -> SOAC rep -> ShowS
[SOAC rep] -> ShowS
SOAC rep -> String
(Int -> SOAC rep -> ShowS)
-> (SOAC rep -> String) -> ([SOAC rep] -> ShowS) -> Show (SOAC rep)
forall rep. RepTypes rep => Int -> SOAC rep -> ShowS
forall rep. RepTypes rep => [SOAC rep] -> ShowS
forall rep. RepTypes rep => SOAC rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> SOAC rep -> ShowS
showsPrec :: Int -> SOAC rep -> ShowS
$cshow :: forall rep. RepTypes rep => SOAC rep -> String
show :: SOAC rep -> String
$cshowList :: forall rep. RepTypes rep => [SOAC rep] -> ShowS
showList :: [SOAC rep] -> ShowS
Show)

-- | Information about computing a single histogram.
data HistOp rep = HistOp
  { forall rep. HistOp rep -> Shape
histShape :: Shape,
    -- | Race factor @RF@ means that only @1/RF@
    -- bins are used.
    forall rep. HistOp rep -> SubExp
histRaceFactor :: SubExp,
    forall rep. HistOp rep -> [VName]
histDest :: [VName],
    forall rep. HistOp rep -> [SubExp]
histNeutral :: [SubExp],
    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 essential parts of a 'Screma' factored out (everything
-- except the input arrays).
data ScremaForm rep = ScremaForm
  { -- | The "main" lambda of the Screma. For a map, this is
    -- equivalent to 'isMapSOAC'. Note that the meaning of the return
    -- value of this lambda depends crucially on exactly which Screma
    -- this is. The parameters will correspond exactly to elements of
    -- the input arrays, however.
    forall rep. ScremaForm rep -> Lambda rep
scremaLambda :: Lambda rep,
    forall rep. ScremaForm rep -> [Scan rep]
scremaScans :: [Scan rep],
    forall rep. ScremaForm rep -> [Reduce rep]
scremaReduces :: [Reduce rep],
    forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda :: Lambda rep
  }
  deriving (ScremaForm rep -> ScremaForm rep -> Bool
(ScremaForm rep -> ScremaForm rep -> Bool)
-> (ScremaForm rep -> ScremaForm rep -> Bool)
-> Eq (ScremaForm rep)
forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
== :: ScremaForm rep -> ScremaForm rep -> Bool
$c/= :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
/= :: ScremaForm rep -> ScremaForm rep -> Bool
Eq, Eq (ScremaForm rep)
Eq (ScremaForm rep) =>
(ScremaForm rep -> ScremaForm rep -> Ordering)
-> (ScremaForm rep -> ScremaForm rep -> Bool)
-> (ScremaForm rep -> ScremaForm rep -> Bool)
-> (ScremaForm rep -> ScremaForm rep -> Bool)
-> (ScremaForm rep -> ScremaForm rep -> Bool)
-> (ScremaForm rep -> ScremaForm rep -> ScremaForm rep)
-> (ScremaForm rep -> ScremaForm rep -> ScremaForm rep)
-> Ord (ScremaForm rep)
ScremaForm rep -> ScremaForm rep -> Bool
ScremaForm rep -> ScremaForm rep -> Ordering
ScremaForm rep -> ScremaForm rep -> ScremaForm 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 (ScremaForm rep)
forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Ordering
forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> ScremaForm rep
$ccompare :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Ordering
compare :: ScremaForm rep -> ScremaForm rep -> Ordering
$c< :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
< :: ScremaForm rep -> ScremaForm rep -> Bool
$c<= :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
<= :: ScremaForm rep -> ScremaForm rep -> Bool
$c> :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
> :: ScremaForm rep -> ScremaForm rep -> Bool
$c>= :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> Bool
>= :: ScremaForm rep -> ScremaForm rep -> Bool
$cmax :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> ScremaForm rep
max :: ScremaForm rep -> ScremaForm rep -> ScremaForm rep
$cmin :: forall rep.
RepTypes rep =>
ScremaForm rep -> ScremaForm rep -> ScremaForm rep
min :: ScremaForm rep -> ScremaForm rep -> ScremaForm rep
Ord, Int -> ScremaForm rep -> ShowS
[ScremaForm rep] -> ShowS
ScremaForm rep -> String
(Int -> ScremaForm rep -> ShowS)
-> (ScremaForm rep -> String)
-> ([ScremaForm rep] -> ShowS)
-> Show (ScremaForm rep)
forall rep. RepTypes rep => Int -> ScremaForm rep -> ShowS
forall rep. RepTypes rep => [ScremaForm rep] -> ShowS
forall rep. RepTypes rep => ScremaForm rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> ScremaForm rep -> ShowS
showsPrec :: Int -> ScremaForm rep -> ShowS
$cshow :: forall rep. RepTypes rep => ScremaForm rep -> String
show :: ScremaForm rep -> String
$cshowList :: forall rep. RepTypes rep => [ScremaForm rep] -> ShowS
showList :: [ScremaForm rep] -> ShowS
Show)

singleBinOp :: (Buildable rep) => [Lambda rep] -> Lambda rep
singleBinOp :: forall rep. Buildable rep => [Lambda rep] -> Lambda rep
singleBinOp [Lambda rep]
lams =
  Lambda
    { lambdaParams :: [LParam rep]
lambdaParams = (Lambda rep -> [Param (TypeBase Shape NoMode)])
-> [Lambda rep] -> [Param (TypeBase Shape NoMode)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Lambda rep -> [Param (TypeBase Shape NoMode)]
Lambda rep -> [LParam rep]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
xParams [Lambda rep]
lams [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a. [a] -> [a] -> [a]
++ (Lambda rep -> [Param (TypeBase Shape NoMode)])
-> [Lambda rep] -> [Param (TypeBase Shape NoMode)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Lambda rep -> [Param (TypeBase Shape NoMode)]
Lambda rep -> [LParam rep]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
yParams [Lambda rep]
lams,
      lambdaReturnType :: [TypeBase Shape NoMode]
lambdaReturnType = (Lambda rep -> [TypeBase Shape NoMode])
-> [Lambda rep] -> [TypeBase Shape NoMode]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType [Lambda rep]
lams,
      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
          ([Stms rep] -> Stms rep
forall a. Monoid a => [a] -> a
mconcat ((Lambda rep -> Stms rep) -> [Lambda rep] -> [Stms rep]
forall a b. (a -> b) -> [a] -> [b]
map (Body rep -> Stms rep
forall rep res. GBody rep res -> Stms rep
bodyStms (Body rep -> Stms rep)
-> (Lambda rep -> Body rep) -> Lambda rep -> Stms 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 -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody) [Lambda rep]
lams))
          ((Lambda rep -> [SubExpRes]) -> [Lambda rep] -> [SubExpRes]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Body rep -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body rep -> [SubExpRes])
-> (Lambda rep -> Body rep) -> Lambda rep -> [SubExpRes]
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 -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody) [Lambda rep]
lams)
    }
  where
    xParams :: GLambda rep t -> [Param (LParamInfo rep)]
xParams GLambda rep t
lam = Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. Int -> [a] -> [a]
take ([t] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (GLambda rep t -> [t]
forall rep t. GLambda rep t -> [t]
lambdaReturnType GLambda rep t
lam)) (GLambda rep t -> [Param (LParamInfo rep)]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams GLambda rep t
lam)
    yParams :: GLambda rep t -> [Param (LParamInfo rep)]
yParams GLambda rep t
lam = Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. Int -> [a] -> [a]
drop ([t] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (GLambda rep t -> [t]
forall rep t. GLambda rep t -> [t]
lambdaReturnType GLambda rep t
lam)) (GLambda rep t -> [Param (LParamInfo rep)]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams GLambda rep t
lam)

-- | How to compute a single scan result.
data Scan rep = Scan
  { forall rep. Scan rep -> Lambda rep
scanLambda :: Lambda rep,
    forall rep. Scan rep -> [SubExp]
scanNeutral :: [SubExp]
  }
  deriving (Scan rep -> Scan rep -> Bool
(Scan rep -> Scan rep -> Bool)
-> (Scan rep -> Scan rep -> Bool) -> Eq (Scan rep)
forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
== :: Scan rep -> Scan rep -> Bool
$c/= :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
/= :: Scan rep -> Scan rep -> Bool
Eq, Eq (Scan rep)
Eq (Scan rep) =>
(Scan rep -> Scan rep -> Ordering)
-> (Scan rep -> Scan rep -> Bool)
-> (Scan rep -> Scan rep -> Bool)
-> (Scan rep -> Scan rep -> Bool)
-> (Scan rep -> Scan rep -> Bool)
-> (Scan rep -> Scan rep -> Scan rep)
-> (Scan rep -> Scan rep -> Scan rep)
-> Ord (Scan rep)
Scan rep -> Scan rep -> Bool
Scan rep -> Scan rep -> Ordering
Scan rep -> Scan rep -> Scan 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 (Scan rep)
forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
forall rep. RepTypes rep => Scan rep -> Scan rep -> Ordering
forall rep. RepTypes rep => Scan rep -> Scan rep -> Scan rep
$ccompare :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Ordering
compare :: Scan rep -> Scan rep -> Ordering
$c< :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
< :: Scan rep -> Scan rep -> Bool
$c<= :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
<= :: Scan rep -> Scan rep -> Bool
$c> :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
> :: Scan rep -> Scan rep -> Bool
$c>= :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
>= :: Scan rep -> Scan rep -> Bool
$cmax :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Scan rep
max :: Scan rep -> Scan rep -> Scan rep
$cmin :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Scan rep
min :: Scan rep -> Scan rep -> Scan rep
Ord, Int -> Scan rep -> ShowS
[Scan rep] -> ShowS
Scan rep -> String
(Int -> Scan rep -> ShowS)
-> (Scan rep -> String) -> ([Scan rep] -> ShowS) -> Show (Scan rep)
forall rep. RepTypes rep => Int -> Scan rep -> ShowS
forall rep. RepTypes rep => [Scan rep] -> ShowS
forall rep. RepTypes rep => Scan rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> Scan rep -> ShowS
showsPrec :: Int -> Scan rep -> ShowS
$cshow :: forall rep. RepTypes rep => Scan rep -> String
show :: Scan rep -> String
$cshowList :: forall rep. RepTypes rep => [Scan rep] -> ShowS
showList :: [Scan rep] -> ShowS
Show)

-- | What are the sizes of reduction results produced by these 'Scan's?
scanSizes :: [Scan rep] -> [Int]
scanSizes :: forall rep. [Scan rep] -> [Int]
scanSizes = (Scan rep -> Int) -> [Scan 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) -> (Scan rep -> [SubExp]) -> Scan 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
. Scan rep -> [SubExp]
forall rep. Scan rep -> [SubExp]
scanNeutral)

-- | How many reduction results are produced by these 'Scan's?
scanResults :: [Scan rep] -> Int
scanResults :: forall rep. [Scan rep] -> Int
scanResults = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([Scan rep] -> [Int]) -> [Scan 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
. [Scan rep] -> [Int]
forall rep. [Scan rep] -> [Int]
scanSizes

-- | Combine multiple scan operators to a single operator.
singleScan :: (Buildable rep) => [Scan rep] -> Scan rep
singleScan :: forall rep. Buildable rep => [Scan rep] -> Scan rep
singleScan [Scan rep]
scans =
  let scan_nes :: [SubExp]
scan_nes = (Scan rep -> [SubExp]) -> [Scan rep] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Scan rep -> [SubExp]
forall rep. Scan rep -> [SubExp]
scanNeutral [Scan rep]
scans
      scan_lam :: Lambda rep
scan_lam = [Lambda rep] -> Lambda rep
forall rep. Buildable rep => [Lambda rep] -> Lambda rep
singleBinOp ([Lambda rep] -> Lambda rep) -> [Lambda rep] -> Lambda rep
forall a b. (a -> b) -> a -> b
$ (Scan rep -> Lambda rep) -> [Scan rep] -> [Lambda rep]
forall a b. (a -> b) -> [a] -> [b]
map Scan rep -> Lambda rep
forall rep. Scan rep -> Lambda rep
scanLambda [Scan rep]
scans
   in Lambda rep -> [SubExp] -> Scan rep
forall rep. Lambda rep -> [SubExp] -> Scan rep
Scan Lambda rep
scan_lam [SubExp]
scan_nes

-- | How to compute a single reduction result.
data Reduce rep = Reduce
  { forall rep. Reduce rep -> Commutativity
redComm :: Commutativity,
    forall rep. Reduce rep -> Lambda rep
redLambda :: Lambda rep,
    forall rep. Reduce rep -> [SubExp]
redNeutral :: [SubExp]
  }
  deriving (Reduce rep -> Reduce rep -> Bool
(Reduce rep -> Reduce rep -> Bool)
-> (Reduce rep -> Reduce rep -> Bool) -> Eq (Reduce rep)
forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
== :: Reduce rep -> Reduce rep -> Bool
$c/= :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
/= :: Reduce rep -> Reduce rep -> Bool
Eq, Eq (Reduce rep)
Eq (Reduce rep) =>
(Reduce rep -> Reduce rep -> Ordering)
-> (Reduce rep -> Reduce rep -> Bool)
-> (Reduce rep -> Reduce rep -> Bool)
-> (Reduce rep -> Reduce rep -> Bool)
-> (Reduce rep -> Reduce rep -> Bool)
-> (Reduce rep -> Reduce rep -> Reduce rep)
-> (Reduce rep -> Reduce rep -> Reduce rep)
-> Ord (Reduce rep)
Reduce rep -> Reduce rep -> Bool
Reduce rep -> Reduce rep -> Ordering
Reduce rep -> Reduce rep -> Reduce 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 (Reduce rep)
forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Ordering
forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Reduce rep
$ccompare :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Ordering
compare :: Reduce rep -> Reduce rep -> Ordering
$c< :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
< :: Reduce rep -> Reduce rep -> Bool
$c<= :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
<= :: Reduce rep -> Reduce rep -> Bool
$c> :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
> :: Reduce rep -> Reduce rep -> Bool
$c>= :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
>= :: Reduce rep -> Reduce rep -> Bool
$cmax :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Reduce rep
max :: Reduce rep -> Reduce rep -> Reduce rep
$cmin :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Reduce rep
min :: Reduce rep -> Reduce rep -> Reduce rep
Ord, Int -> Reduce rep -> ShowS
[Reduce rep] -> ShowS
Reduce rep -> String
(Int -> Reduce rep -> ShowS)
-> (Reduce rep -> String)
-> ([Reduce rep] -> ShowS)
-> Show (Reduce rep)
forall rep. RepTypes rep => Int -> Reduce rep -> ShowS
forall rep. RepTypes rep => [Reduce rep] -> ShowS
forall rep. RepTypes rep => Reduce rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> Reduce rep -> ShowS
showsPrec :: Int -> Reduce rep -> ShowS
$cshow :: forall rep. RepTypes rep => Reduce rep -> String
show :: Reduce rep -> String
$cshowList :: forall rep. RepTypes rep => [Reduce rep] -> ShowS
showList :: [Reduce rep] -> ShowS
Show)

-- | What are the sizes of reduction results produced by these 'Reduce's?
redSizes :: [Reduce rep] -> [Int]
redSizes :: forall rep. [Reduce rep] -> [Int]
redSizes = (Reduce rep -> Int) -> [Reduce 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) -> (Reduce rep -> [SubExp]) -> Reduce 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
. Reduce rep -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral)

-- | How many reduction results are produced by these 'Reduce's?
redResults :: [Reduce rep] -> Int
redResults :: forall rep. [Reduce rep] -> Int
redResults = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([Reduce rep] -> [Int]) -> [Reduce 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
. [Reduce rep] -> [Int]
forall rep. [Reduce rep] -> [Int]
redSizes

-- | Combine multiple reduction operators to a single operator.
singleReduce :: (Buildable rep) => [Reduce rep] -> Reduce rep
singleReduce :: forall rep. Buildable rep => [Reduce rep] -> Reduce rep
singleReduce [Reduce rep]
reds =
  let red_nes :: [SubExp]
red_nes = (Reduce rep -> [SubExp]) -> [Reduce rep] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Reduce rep -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral [Reduce rep]
reds
      red_lam :: Lambda rep
red_lam = [Lambda rep] -> Lambda rep
forall rep. Buildable rep => [Lambda rep] -> Lambda rep
singleBinOp ([Lambda rep] -> Lambda rep) -> [Lambda rep] -> Lambda rep
forall a b. (a -> b) -> a -> b
$ (Reduce rep -> Lambda rep) -> [Reduce rep] -> [Lambda rep]
forall a b. (a -> b) -> [a] -> [b]
map Reduce rep -> Lambda rep
forall rep. Reduce rep -> Lambda rep
redLambda [Reduce rep]
reds
   in Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce ([Commutativity] -> Commutativity
forall a. Monoid a => [a] -> a
mconcat ((Reduce rep -> Commutativity) -> [Reduce rep] -> [Commutativity]
forall a b. (a -> b) -> [a] -> [b]
map Reduce rep -> Commutativity
forall rep. Reduce rep -> Commutativity
redComm [Reduce rep]
reds)) Lambda rep
red_lam [SubExp]
red_nes

-- | The types produced by a single 'Screma', given the size of the
-- input array.
scremaType :: SubExp -> ScremaForm rep -> [Type]
scremaType :: forall rep. SubExp -> ScremaForm rep -> [TypeBase Shape NoMode]
scremaType SubExp
w (ScremaForm Lambda rep
_map_lam [Scan rep]
_scans [Reduce rep]
reds Lambda rep
post_lam) =
  [TypeBase Shape NoMode]
red_tps [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Semigroup a => a -> a -> a
<> (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (TypeBase Shape NoMode -> SubExp -> TypeBase Shape NoMode
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
`arrayOfRow` SubExp
w) (Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
post_lam)
  where
    red_tps :: [TypeBase Shape NoMode]
red_tps = (Reduce rep -> [TypeBase Shape NoMode])
-> [Reduce rep] -> [TypeBase Shape NoMode]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (Lambda rep -> [TypeBase Shape NoMode])
-> (Reduce rep -> Lambda rep)
-> Reduce 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
. Reduce rep -> Lambda rep
forall rep. Reduce rep -> Lambda rep
redLambda) [Reduce rep]
reds

-- | Creates let-bindings to compose two lambda functions (producer →
-- consumer).
--
-- When composing two operations where the outputs of one lambda
-- (producer) flow into the inputs of another lambda (consumer), this
-- function generates the necessary let-bindings that connect matching
-- outputs to inputs.
--
-- The function looks for outputs from the producer that correspond to
-- inputs expected by the consumer, and creates bindings like: let
-- consumer_param = producer_result
--
-- This allows the two lambdas to be composed into a operation.
--
-- Returns: Statements containing let-bindings for each matched
-- producer output → consumer input pair. Preserves certificates from
-- the producer results.
--
-- Example: If out_p[i] == inp_c[j], then producer's result[i] is
-- bound to consumer's parameter[j]. Unmatched outputs (producer
-- results not consumed) are omitted.
composeBinds ::
  (Buildable rep, Ord a) =>
  -- | Producer lambda
  Lambda rep ->
  -- | Producer outputs to match
  [a] ->
  -- | Consumer inputs to match
  [a] ->
  -- | Consumer lambda
  Lambda rep ->
  -- | Let-bindings connecting them
  Stms rep
composeBinds :: forall rep a.
(Buildable rep, Ord a) =>
Lambda rep -> [a] -> [a] -> Lambda rep -> Stms rep
composeBinds Lambda rep
lam_p [a]
out_p [a]
inp_c Lambda rep
lam_c =
  [Stm rep] -> Stms rep
forall rep. [Stm rep] -> Stms rep
stmsFromList ([Stm rep] -> Stms rep)
-> ([(a, SubExpRes, TypeBase Shape NoMode)] -> [Stm rep])
-> [(a, SubExpRes, TypeBase Shape NoMode)]
-> Stms 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
. ((a, SubExpRes, TypeBase Shape NoMode) -> Maybe (Stm rep))
-> [(a, SubExpRes, TypeBase Shape NoMode)] -> [Stm rep]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (a, SubExpRes, TypeBase Shape NoMode) -> Maybe (Stm rep)
bindResToPar ([(a, SubExpRes, TypeBase Shape NoMode)] -> Stms rep)
-> [(a, SubExpRes, TypeBase Shape NoMode)] -> Stms rep
forall a b. (a -> b) -> a -> b
$ [a]
-> [SubExpRes]
-> [TypeBase Shape NoMode]
-> [(a, SubExpRes, TypeBase Shape NoMode)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [a]
out_p [SubExpRes]
res_p [TypeBase Shape NoMode]
ts_p
  where
    ts_p :: [TypeBase Shape NoMode]
ts_p = Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam_p
    res_p :: [SubExpRes]
res_p = GBody rep SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody rep SubExpRes -> [SubExpRes])
-> GBody rep SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam_p

    inp_c_map :: Map a VName
inp_c_map =
      [(a, VName)] -> Map a VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, VName)] -> Map a VName)
-> ([VName] -> [(a, VName)]) -> [VName] -> Map a 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
. [a] -> [VName] -> [(a, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
inp_c ([VName] -> Map a VName) -> [VName] -> Map a VName
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
paramName (Param (TypeBase Shape NoMode) -> VName)
-> [Param (TypeBase Shape NoMode)] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda rep -> [LParam rep]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams Lambda rep
lam_c

    bindResToPar :: (a, SubExpRes, TypeBase Shape NoMode) -> Maybe (Stm rep)
bindResToPar (a
out, SubExpRes
res, TypeBase Shape NoMode
t) =
      case a -> Map a VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
out Map a VName
inp_c_map of
        Just VName
name ->
          Stm rep -> Maybe (Stm rep)
forall a. a -> Maybe a
Just (Stm rep -> Maybe (Stm rep)) -> Stm rep -> Maybe (Stm rep)
forall a b. (a -> b) -> a -> b
$ Certs -> Stm rep -> Stm rep
forall rep. Certs -> Stm rep -> Stm rep
certify Certs
cs (Stm rep -> Stm rep) -> Stm rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [VName -> TypeBase Shape NoMode -> Ident
Ident VName
name TypeBase Shape NoMode
t] (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ 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
e
          where
            SubExpRes Certs
cs SubExp
e = SubExpRes
res
        Maybe VName
Nothing -> Maybe (Stm rep)
forall a. Maybe a
Nothing

-- | Construct a lambda that takes parameters of the given types and
-- simply returns them unchanged.
mkIdentityLambda ::
  (Buildable rep, MonadFreshNames m) =>
  [Type] ->
  m (Lambda rep)
mkIdentityLambda :: forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[TypeBase Shape NoMode] -> m (Lambda rep)
mkIdentityLambda [TypeBase Shape NoMode]
ts = do
  params <- (TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode)))
-> [TypeBase Shape NoMode] -> m [Param (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 (Name -> TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"x") [TypeBase Shape NoMode]
ts
  pure
    Lambda
      { lambdaParams = params,
        lambdaBody = mkBody mempty $ varsRes $ map paramName params,
        lambdaReturnType = ts
      }

-- | A lambda with no parameters that returns no values.
nilFn :: (Buildable rep) => Lambda rep
nilFn :: forall rep. Buildable rep => Lambda rep
nilFn = [LParam rep]
-> [TypeBase Shape NoMode]
-> Body rep
-> GLambda rep (TypeBase Shape NoMode)
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param (TypeBase Shape NoMode)]
[LParam rep]
forall a. Monoid a => a
mempty [TypeBase Shape NoMode]
forall a. Monoid a => a
mempty (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 Stms rep
forall a. Monoid a => a
mempty [SubExpRes]
forall a. Monoid a => a
mempty)

-- | Construct a Screma with possibly multiple scans, and the given
-- map function.
scanomapSOAC ::
  (Buildable rep, MonadFreshNames m) =>
  [Scan rep] ->
  Lambda rep ->
  m (ScremaForm rep)
scanomapSOAC :: forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> Lambda rep -> m (ScremaForm rep)
scanomapSOAC [Scan rep]
scans Lambda rep
lam =
  Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm Lambda rep
lam [Scan rep]
scans [] (Lambda rep -> ScremaForm rep)
-> m (Lambda rep) -> m (ScremaForm rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TypeBase Shape NoMode] -> m (Lambda rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[TypeBase Shape NoMode] -> m (Lambda rep)
mkIdentityLambda (Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam)

-- | Construct a Screma with possibly multiple scans,
-- the given map function, and a given post lambda.
maposcanomapSOAC ::
  (Buildable rep, MonadFreshNames m) =>
  Lambda rep ->
  [Scan rep] ->
  Lambda rep ->
  m (ScremaForm rep)
maposcanomapSOAC :: forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
Lambda rep -> [Scan rep] -> Lambda rep -> m (ScremaForm rep)
maposcanomapSOAC Lambda rep
pre_lam [] Lambda rep
post_lam = do
  new_post_lam <- [TypeBase Shape NoMode] -> m (Lambda rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[TypeBase Shape NoMode] -> m (Lambda rep)
mkIdentityLambda ([TypeBase Shape NoMode] -> m (Lambda rep))
-> [TypeBase Shape NoMode] -> m (Lambda rep)
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
post_lam
  let new_pre_lam =
        Lambda
          { lambdaParams :: [LParam rep]
lambdaParams = Lambda rep -> [LParam rep]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams Lambda rep
pre_lam,
            lambdaReturnType :: [TypeBase Shape NoMode]
lambdaReturnType = Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
post_lam,
            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 Stms rep
new_stms [SubExpRes]
new_res
          }
  pure $ ScremaForm new_pre_lam [] [] new_post_lam
  where
    new_res :: [SubExpRes]
new_res = Body rep -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body rep -> [SubExpRes]) -> Body rep -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
post_lam
    stmsFromLam :: GLambda rep t -> Stms rep
stmsFromLam = GBody rep SubExpRes -> Stms rep
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody rep SubExpRes -> Stms rep)
-> (GLambda rep t -> GBody rep SubExpRes)
-> GLambda rep t
-> Stms 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
. GLambda rep t -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody
    deps :: [Int]
deps = [Int
0 .. [TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([TypeBase Shape NoMode] -> Int) -> [TypeBase Shape NoMode] -> Int
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
pre_lam]
    new_stms :: Stms rep
new_stms =
      Lambda rep -> Stms rep
forall {rep} {t}. GLambda rep t -> Stms rep
stmsFromLam Lambda rep
pre_lam
        Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> [Int] -> [Int] -> Lambda rep -> Stms rep
forall rep a.
(Buildable rep, Ord a) =>
Lambda rep -> [a] -> [a] -> Lambda rep -> Stms rep
composeBinds Lambda rep
pre_lam [Int]
deps [Int]
deps Lambda rep
post_lam
        Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> Stms rep
forall {rep} {t}. GLambda rep t -> Stms rep
stmsFromLam Lambda rep
post_lam
maposcanomapSOAC Lambda rep
lam [Scan rep]
scans Lambda rep
post_lam =
  ScremaForm rep -> m (ScremaForm rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ScremaForm rep -> m (ScremaForm rep))
-> ScremaForm rep -> m (ScremaForm rep)
forall a b. (a -> b) -> a -> b
$ Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm Lambda rep
lam [Scan rep]
scans [] Lambda rep
post_lam

-- | Construct a Screma with possibly multiple reductions, and
-- the given map function.
redomapSOAC ::
  (Buildable rep, MonadFreshNames m) =>
  [Reduce rep] ->
  Lambda rep ->
  m (ScremaForm rep)
redomapSOAC :: forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> Lambda rep -> m (ScremaForm rep)
redomapSOAC [Reduce rep]
reds Lambda rep
lam = Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm Lambda rep
lam [] [Reduce rep]
reds (Lambda rep -> ScremaForm rep)
-> m (Lambda rep) -> m (ScremaForm rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TypeBase Shape NoMode] -> m (Lambda rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[TypeBase Shape NoMode] -> m (Lambda rep)
mkIdentityLambda [TypeBase Shape NoMode]
map_ts
  where
    map_ts :: [TypeBase Shape NoMode]
map_ts = Int -> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Int -> [a] -> [a]
drop ([Reduce rep] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce rep]
reds) ([TypeBase Shape NoMode] -> [TypeBase Shape NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam

-- | Construct a Screma with possibly multiple scans, and identity map
-- function.
scanSOAC ::
  (Buildable rep, MonadFreshNames m) =>
  [Scan rep] ->
  m (ScremaForm rep)
scanSOAC :: forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
scanSOAC [Scan rep]
scans = [Scan rep] -> Lambda rep -> m (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> Lambda rep -> m (ScremaForm rep)
scanomapSOAC [Scan rep]
scans (Lambda rep -> m (ScremaForm rep))
-> m (Lambda rep) -> m (ScremaForm rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [TypeBase Shape NoMode] -> m (Lambda rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[TypeBase Shape NoMode] -> m (Lambda rep)
mkIdentityLambda [TypeBase Shape NoMode]
ts
  where
    ts :: [TypeBase Shape NoMode]
ts = (Scan rep -> [TypeBase Shape NoMode])
-> [Scan rep] -> [TypeBase Shape NoMode]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (Lambda rep -> [TypeBase Shape NoMode])
-> (Scan rep -> Lambda rep) -> Scan 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
. Scan rep -> Lambda rep
forall rep. Scan rep -> Lambda rep
scanLambda) [Scan rep]
scans

-- | Construct a Screma with possibly multiple reductions, and
-- identity map function.
reduceSOAC ::
  (Buildable rep, MonadFreshNames m) =>
  [Reduce rep] ->
  m (ScremaForm rep)
reduceSOAC :: forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
reduceSOAC [Reduce rep]
reds = [Reduce rep] -> Lambda rep -> m (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> Lambda rep -> m (ScremaForm rep)
redomapSOAC [Reduce rep]
reds (Lambda rep -> m (ScremaForm rep))
-> m (Lambda rep) -> m (ScremaForm rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [TypeBase Shape NoMode] -> m (Lambda rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[TypeBase Shape NoMode] -> m (Lambda rep)
mkIdentityLambda [TypeBase Shape NoMode]
ts
  where
    ts :: [TypeBase Shape NoMode]
ts = (Reduce rep -> [TypeBase Shape NoMode])
-> [Reduce rep] -> [TypeBase Shape NoMode]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (Lambda rep -> [TypeBase Shape NoMode])
-> (Reduce rep -> Lambda rep)
-> Reduce 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
. Reduce rep -> Lambda rep
forall rep. Reduce rep -> Lambda rep
redLambda) [Reduce rep]
reds

-- | Construct a Screma corresponding to a map.
mapSOAC ::
  (Buildable rep, MonadFreshNames m) =>
  Lambda rep ->
  m (ScremaForm rep)
mapSOAC :: forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
Lambda rep -> m (ScremaForm rep)
mapSOAC Lambda rep
lam = do
  post_lam <- [TypeBase Shape NoMode] -> m (Lambda rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[TypeBase Shape NoMode] -> m (Lambda rep)
mkIdentityLambda ([TypeBase Shape NoMode] -> m (Lambda rep))
-> [TypeBase Shape NoMode] -> m (Lambda rep)
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam
  pure $ ScremaForm lam [] [] post_lam

-- | Does this Screma correspond to a scan-map composition?
isScanomapSOAC :: ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
isScanomapSOAC :: forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
isScanomapSOAC (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam) = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [Reduce rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Reduce rep]
reds
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Scan rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Scan rep]
scans
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Bool
forall rep. Lambda rep -> Bool
isIdentityLambda Lambda rep
post_lam
  ([Scan rep], Lambda rep) -> Maybe ([Scan rep], Lambda rep)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Scan rep]
scans, Lambda rep
map_lam)

isMaposcanomapSOAC :: ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC :: forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam) = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [Reduce rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Reduce rep]
reds
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Scan rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Scan rep]
scans
  (Lambda rep, [Scan rep], Lambda rep)
-> Maybe (Lambda rep, [Scan rep], Lambda rep)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda rep
post_lam, [Scan rep]
scans, Lambda rep
map_lam)

-- | Does this Screma correspond to pure scan?
isScanSOAC :: ScremaForm rep -> Maybe [Scan rep]
isScanSOAC :: forall rep. ScremaForm rep -> Maybe [Scan rep]
isScanSOAC ScremaForm rep
form = do
  (scans, map_lam) <- ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
isScanomapSOAC ScremaForm rep
form
  guard $ isIdentityLambda map_lam
  guard $ length (lambdaReturnType map_lam) == scanResults scans
  pure scans

-- | Does this Screma correspond to a reduce-map composition?
isRedomapSOAC :: ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC :: forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam) = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [Scan rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Scan rep]
scans
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Reduce rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Reduce rep]
reds
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Bool
forall rep. Lambda rep -> Bool
isIdentityLambda Lambda rep
post_lam
  ([Reduce rep], Lambda rep) -> Maybe ([Reduce rep], Lambda rep)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Reduce rep]
reds, Lambda rep
map_lam)

-- | Does this Screma correspond to a pure reduce?
isReduceSOAC :: ScremaForm rep -> Maybe [Reduce rep]
isReduceSOAC :: forall rep. ScremaForm rep -> Maybe [Reduce rep]
isReduceSOAC ScremaForm rep
form = do
  (reds, map_lam) <- ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm rep
form
  guard $ isIdentityLambda map_lam
  guard $ length (lambdaReturnType map_lam) == redResults reds
  pure reds

-- | Does this Screma correspond to a simple map, without any
-- reduction or scan results?
isMapSOAC :: ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC :: forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam) = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [Scan rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Scan rep]
scans
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [Reduce rep] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Reduce rep]
reds
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Bool
forall rep. Lambda rep -> Bool
isIdentityLambda Lambda rep
post_lam
  Lambda rep -> Maybe (Lambda rep)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda rep
map_lam

-- | Like 'Mapper', but just for 'SOAC's.
data SOACMapper frep trep m = SOACMapper
  { forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp :: SubExp -> m SubExp,
    forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda :: Lambda frep -> m (Lambda trep),
    forall frep trep (m :: * -> *).
SOACMapper frep trep m -> ExtLambda frep -> m (ExtLambda trep)
mapOnSOACExtLambda :: ExtLambda frep -> m (ExtLambda trep),
    forall frep trep (m :: * -> *).
SOACMapper frep trep m -> VName -> m VName
mapOnSOACVName :: VName -> m VName
  }

-- | A mapper that simply returns the SOAC verbatim.
identitySOACMapper :: forall rep m. (Monad m) => SOACMapper rep rep m
identitySOACMapper :: forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper =
  SOACMapper
    { mapOnSOACSubExp :: SubExp -> m SubExp
mapOnSOACSubExp = SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSOACLambda :: Lambda rep -> m (Lambda rep)
mapOnSOACLambda = Lambda rep -> m (Lambda rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSOACExtLambda :: ExtLambda rep -> m (ExtLambda rep)
mapOnSOACExtLambda = ExtLambda rep -> m (ExtLambda rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnSOACVName :: VName -> m VName
mapOnSOACVName = VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    }

-- | Map a monadic action across the immediate children of a
-- SOAC.  The mapping does not descend recursively into subexpressions
-- and is done left-to-right.
mapSOACM ::
  (Monad m) =>
  SOACMapper frep trep m ->
  SOAC frep ->
  m (SOAC trep)
mapSOACM :: forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper frep trep m
tv (JVP Shape
shape [SubExp]
args [SubExp]
vec Lambda frep
lam) =
  Shape -> [SubExp] -> [SubExp] -> Lambda trep -> SOAC trep
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
JVP
    (Shape -> [SubExp] -> [SubExp] -> Lambda trep -> SOAC trep)
-> m Shape -> m ([SubExp] -> [SubExp] -> Lambda trep -> SOAC 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) Shape
shape
    m ([SubExp] -> [SubExp] -> Lambda trep -> SOAC trep)
-> m [SubExp] -> m ([SubExp] -> Lambda trep -> SOAC 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
args
    m ([SubExp] -> Lambda trep -> SOAC trep)
-> m [SubExp] -> m (Lambda trep -> SOAC 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
vec
    m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
lam
mapSOACM SOACMapper frep trep m
tv (VJP Shape
shape [SubExp]
args [SubExp]
vec Lambda frep
lam) =
  Shape -> [SubExp] -> [SubExp] -> Lambda trep -> SOAC trep
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
VJP
    (Shape -> [SubExp] -> [SubExp] -> Lambda trep -> SOAC trep)
-> m Shape -> m ([SubExp] -> [SubExp] -> Lambda trep -> SOAC 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) Shape
shape
    m ([SubExp] -> [SubExp] -> Lambda trep -> SOAC trep)
-> m [SubExp] -> m ([SubExp] -> Lambda trep -> SOAC 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
args
    m ([SubExp] -> Lambda trep -> SOAC trep)
-> m [SubExp] -> m (Lambda trep -> SOAC 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
vec
    m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
lam
mapSOACM SOACMapper frep trep m
tv (WithVJP [SubExp]
args Lambda frep
lam0 Lambda frep
lam1) =
  [SubExp] -> Lambda trep -> Lambda trep -> SOAC trep
forall rep. [SubExp] -> Lambda rep -> Lambda rep -> SOAC rep
WithVJP
    ([SubExp] -> Lambda trep -> Lambda trep -> SOAC trep)
-> m [SubExp] -> m (Lambda trep -> Lambda trep -> SOAC trep)
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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
args
    m (Lambda trep -> Lambda trep -> SOAC trep)
-> m (Lambda trep) -> m (Lambda trep -> SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
lam0
    m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
lam1
mapSOACM SOACMapper frep trep m
tv (Stream SubExp
size [VName]
arrs [SubExp]
accs Lambda frep
lam) =
  SubExp -> [VName] -> [SubExp] -> Lambda trep -> SOAC trep
forall rep. SubExp -> [VName] -> [SubExp] -> Lambda rep -> SOAC rep
Stream
    (SubExp -> [VName] -> [SubExp] -> Lambda trep -> SOAC trep)
-> m SubExp -> m ([VName] -> [SubExp] -> Lambda trep -> SOAC trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv SubExp
size
    m ([VName] -> [SubExp] -> Lambda trep -> SOAC trep)
-> m [VName] -> m ([SubExp] -> Lambda trep -> SOAC 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 (SOACMapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> VName -> m VName
mapOnSOACVName SOACMapper frep trep m
tv) [VName]
arrs
    m ([SubExp] -> Lambda trep -> SOAC trep)
-> m [SubExp] -> m (Lambda trep -> SOAC 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
accs
    m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
lam
mapSOACM SOACMapper frep trep m
tv (Hist SubExp
w [VName]
arrs [HistOp frep]
ops Lambda frep
bucket_fun) =
  SubExp -> [VName] -> [HistOp trep] -> Lambda trep -> SOAC trep
forall rep.
SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
Hist
    (SubExp -> [VName] -> [HistOp trep] -> Lambda trep -> SOAC trep)
-> m SubExp
-> m ([VName] -> [HistOp trep] -> Lambda trep -> SOAC trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv SubExp
w
    m ([VName] -> [HistOp trep] -> Lambda trep -> SOAC trep)
-> m [VName] -> m ([HistOp trep] -> Lambda trep -> SOAC 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 (SOACMapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> VName -> m VName
mapOnSOACVName SOACMapper frep trep m
tv) [VName]
arrs
    m ([HistOp trep] -> Lambda trep -> SOAC trep)
-> m [HistOp trep] -> m (Lambda trep -> SOAC 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 Shape
shape SubExp
rf [VName]
op_arrs [SubExp]
nes Lambda frep
op) ->
          Shape
-> SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep
forall rep.
Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
HistOp
            (Shape
 -> SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep)
-> m Shape
-> m (SubExp -> [VName] -> [SubExp] -> 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) Shape
shape
            m (SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep)
-> m SubExp
-> m ([VName] -> [SubExp] -> 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
<*> SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv SubExp
rf
            m ([VName] -> [SubExp] -> Lambda trep -> HistOp trep)
-> m [VName] -> m ([SubExp] -> 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 (SOACMapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> VName -> m VName
mapOnSOACVName SOACMapper frep trep m
tv) [VName]
op_arrs
            m ([SubExp] -> Lambda trep -> HistOp trep)
-> m [SubExp] -> 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 -> 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
nes
            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
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
op
      )
      [HistOp frep]
ops
    m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
bucket_fun
mapSOACM SOACMapper frep trep m
tv (Screma SubExp
w [VName]
arrs (ScremaForm Lambda frep
map_lam [Scan frep]
scans [Reduce frep]
reds Lambda frep
post_lam)) =
  SubExp -> [VName] -> ScremaForm trep -> SOAC trep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma
    (SubExp -> [VName] -> ScremaForm trep -> SOAC trep)
-> m SubExp -> m ([VName] -> ScremaForm trep -> SOAC trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv SubExp
w
    m ([VName] -> ScremaForm trep -> SOAC trep)
-> m [VName] -> m (ScremaForm trep -> SOAC 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 (SOACMapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> VName -> m VName
mapOnSOACVName SOACMapper frep trep m
tv) [VName]
arrs
    m (ScremaForm trep -> SOAC trep)
-> m (ScremaForm trep) -> m (SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ( Lambda trep
-> [Scan trep] -> [Reduce trep] -> Lambda trep -> ScremaForm trep
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm
            (Lambda trep
 -> [Scan trep] -> [Reduce trep] -> Lambda trep -> ScremaForm trep)
-> m (Lambda trep)
-> m ([Scan trep]
      -> [Reduce trep] -> Lambda trep -> ScremaForm trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
map_lam
            m ([Scan trep] -> [Reduce trep] -> Lambda trep -> ScremaForm trep)
-> m [Scan trep]
-> m ([Reduce trep] -> Lambda trep -> ScremaForm trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Scan frep -> m (Scan trep)) -> [Scan frep] -> m [Scan 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 (SOACMapper frep trep m -> Scan frep -> m (Scan trep)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> Scan frep -> m (Scan trep)
mapOnSOACScan SOACMapper frep trep m
tv) [Scan frep]
scans
            m ([Reduce trep] -> Lambda trep -> ScremaForm trep)
-> m [Reduce trep] -> m (Lambda trep -> ScremaForm trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Reduce frep -> m (Reduce trep))
-> [Reduce frep] -> m [Reduce 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 (SOACMapper frep trep m -> Reduce frep -> m (Reduce trep)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> Reduce frep -> m (Reduce trep)
mapOnSOACReduce SOACMapper frep trep m
tv) [Reduce frep]
reds
            m (Lambda trep -> ScremaForm trep)
-> m (Lambda trep) -> m (ScremaForm trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
post_lam
        )
mapSOACM SOACMapper frep trep m
tv (FlatMap SubExp
w [VName]
arrs ExtLambda frep
lam) =
  SubExp -> [VName] -> ExtLambda trep -> SOAC trep
forall rep. SubExp -> [VName] -> ExtLambda rep -> SOAC rep
FlatMap
    (SubExp -> [VName] -> ExtLambda trep -> SOAC trep)
-> m SubExp -> m ([VName] -> ExtLambda trep -> SOAC trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv SubExp
w
    m ([VName] -> ExtLambda trep -> SOAC trep)
-> m [VName] -> m (ExtLambda trep -> SOAC 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 (SOACMapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> VName -> m VName
mapOnSOACVName SOACMapper frep trep m
tv) [VName]
arrs
    m (ExtLambda trep -> SOAC trep)
-> m (ExtLambda trep) -> m (SOAC trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SOACMapper frep trep m -> ExtLambda frep -> m (ExtLambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> ExtLambda frep -> m (ExtLambda trep)
mapOnSOACExtLambda SOACMapper frep trep m
tv ExtLambda frep
lam

mapOnSOACScan :: (Monad m) => SOACMapper frep trep m -> Scan frep -> m (Scan trep)
mapOnSOACScan :: forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> Scan frep -> m (Scan trep)
mapOnSOACScan SOACMapper frep trep m
tv (Scan Lambda frep
red_lam [SubExp]
red_nes) =
  Lambda trep -> [SubExp] -> Scan trep
forall rep. Lambda rep -> [SubExp] -> Scan rep
Scan
    (Lambda trep -> [SubExp] -> Scan trep)
-> m (Lambda trep) -> m ([SubExp] -> Scan trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
red_lam
    m ([SubExp] -> Scan trep) -> m [SubExp] -> m (Scan 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
red_nes

mapOnSOACReduce :: (Monad m) => SOACMapper frep trep m -> Reduce frep -> m (Reduce trep)
mapOnSOACReduce :: forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> Reduce frep -> m (Reduce trep)
mapOnSOACReduce SOACMapper frep trep m
tv (Reduce Commutativity
comm Lambda frep
red_lam [SubExp]
red_nes) =
  Commutativity -> Lambda trep -> [SubExp] -> Reduce trep
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
comm
    (Lambda trep -> [SubExp] -> Reduce trep)
-> m (Lambda trep) -> m ([SubExp] -> Reduce trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnSOACLambda SOACMapper frep trep m
tv Lambda frep
red_lam
    m ([SubExp] -> Reduce trep) -> m [SubExp] -> m (Reduce 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 (SOACMapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
SOACMapper frep trep m -> SubExp -> m SubExp
mapOnSOACSubExp SOACMapper frep trep m
tv) [SubExp]
red_nes

-- | A helper for defining 'TraverseOpStms'.
traverseSOACStms :: (Monad m) => OpStmsTraverser m (SOAC rep) rep
traverseSOACStms :: forall (m :: * -> *) rep.
Monad m =>
OpStmsTraverser m (SOAC rep) rep
traverseSOACStms Scope rep -> Stms rep -> m (Stms rep)
f = SOACMapper rep rep m -> SOAC rep -> m (SOAC rep)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper rep rep m
mapper
  where
    mapper :: SOACMapper rep rep m
mapper =
      SOACMapper (ZonkAny 3) (ZonkAny 3) m
forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper
        { mapOnSOACLambda = traverseLambdaStms f,
          mapOnSOACExtLambda = traverseLambdaStms f
        }

instance (ASTRep rep) => FreeIn (Scan rep) where
  freeIn' :: Scan rep -> FV
freeIn' (Scan Lambda rep
lam [SubExp]
ne) = Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Lambda rep
lam FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> FV
forall a. FreeIn a => a -> FV
freeIn' [SubExp]
ne

instance (ASTRep rep) => FreeIn (Reduce rep) where
  freeIn' :: Reduce rep -> FV
freeIn' (Reduce Commutativity
_ Lambda rep
lam [SubExp]
ne) = Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Lambda rep
lam FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> FV
forall a. FreeIn a => a -> FV
freeIn' [SubExp]
ne

instance (ASTRep rep) => FreeIn (ScremaForm rep) where
  freeIn' :: ScremaForm rep -> FV
freeIn' (ScremaForm Lambda rep
scans [Scan rep]
reds [Reduce rep]
lam Lambda rep
post_lam) =
    Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Lambda rep
scans FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [Scan rep] -> FV
forall a. FreeIn a => a -> FV
freeIn' [Scan rep]
reds FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [Reduce rep] -> FV
forall a. FreeIn a => a -> FV
freeIn' [Reduce rep]
lam FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Lambda rep
post_lam

instance (ASTRep rep) => FreeIn (HistOp rep) where
  freeIn' :: HistOp rep -> FV
freeIn' (HistOp Shape
w SubExp
rf [VName]
dests [SubExp]
nes Lambda rep
lam) =
    Shape -> FV
forall a. FreeIn a => a -> FV
freeIn' Shape
w FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn' SubExp
rf FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [VName] -> FV
forall a. FreeIn a => a -> FV
freeIn' [VName]
dests FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> FV
forall a. FreeIn a => a -> FV
freeIn' [SubExp]
nes FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Lambda rep
lam

instance (ASTRep rep) => FreeIn (SOAC rep) where
  freeIn' :: SOAC rep -> FV
freeIn' = (State FV (SOAC rep) -> FV -> FV)
-> FV -> State FV (SOAC rep) -> FV
forall a b c. (a -> b -> c) -> b -> a -> c
flip State FV (SOAC rep) -> FV -> FV
forall s a. State s a -> s -> s
execState FV
forall a. Monoid a => a
mempty (State FV (SOAC rep) -> FV)
-> (SOAC rep -> State FV (SOAC rep)) -> SOAC rep -> FV
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
. SOACMapper rep rep (StateT FV Identity)
-> SOAC rep -> State FV (SOAC rep)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper rep rep (StateT FV Identity)
free
    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 :: SOACMapper rep rep (StateT FV Identity)
free =
        SOACMapper
          { mapOnSOACSubExp :: SubExp -> StateT FV Identity SubExp
mapOnSOACSubExp = (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',
            mapOnSOACLambda :: Lambda rep -> StateT FV Identity (Lambda rep)
mapOnSOACLambda = (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',
            mapOnSOACExtLambda :: ExtLambda rep -> StateT FV Identity (ExtLambda rep)
mapOnSOACExtLambda = (ExtLambda rep -> FV)
-> ExtLambda rep -> StateT FV Identity (ExtLambda rep)
forall {m :: * -> *} {s} {b}.
(MonadState s m, Semigroup s) =>
(b -> s) -> b -> m b
walk ExtLambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
            mapOnSOACVName :: VName -> StateT FV Identity VName
mapOnSOACVName = (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'
          }

instance (ASTRep rep) => Substitute (SOAC rep) where
  substituteNames :: Map VName VName -> SOAC rep -> SOAC rep
substituteNames Map VName VName
subst =
    Identity (SOAC rep) -> SOAC rep
forall a. Identity a -> a
runIdentity (Identity (SOAC rep) -> SOAC rep)
-> (SOAC rep -> Identity (SOAC rep)) -> SOAC rep -> SOAC 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
. SOACMapper rep rep Identity -> SOAC rep -> Identity (SOAC rep)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper rep rep Identity
substitute
    where
      substitute :: SOACMapper rep rep Identity
substitute =
        SOACMapper
          { mapOnSOACSubExp :: SubExp -> Identity SubExp
mapOnSOACSubExp = 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,
            mapOnSOACLambda :: Lambda rep -> Identity (Lambda rep)
mapOnSOACLambda = 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,
            mapOnSOACExtLambda :: ExtLambda rep -> Identity (ExtLambda rep)
mapOnSOACExtLambda = ExtLambda rep -> Identity (ExtLambda rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExtLambda rep -> Identity (ExtLambda rep))
-> (ExtLambda rep -> ExtLambda rep)
-> ExtLambda rep
-> Identity (ExtLambda 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 -> ExtLambda rep -> ExtLambda rep
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst,
            mapOnSOACVName :: VName -> Identity VName
mapOnSOACVName = 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
          }

instance (ASTRep rep) => Rename (SOAC rep) where
  rename :: SOAC rep -> RenameM (SOAC rep)
rename = SOACMapper rep rep RenameM -> SOAC rep -> RenameM (SOAC rep)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper rep rep RenameM
renamer
    where
      renamer :: SOACMapper rep rep RenameM
renamer = (SubExp -> RenameM SubExp)
-> (Lambda rep -> RenameM (Lambda rep))
-> (ExtLambda rep -> RenameM (ExtLambda rep))
-> (VName -> RenameM VName)
-> SOACMapper rep rep RenameM
forall frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Lambda frep -> m (Lambda trep))
-> (ExtLambda frep -> m (ExtLambda trep))
-> (VName -> m VName)
-> SOACMapper frep trep m
SOACMapper SubExp -> RenameM SubExp
forall a. Rename a => a -> RenameM a
rename Lambda rep -> RenameM (Lambda rep)
forall a. Rename a => a -> RenameM a
rename ExtLambda rep -> RenameM (ExtLambda rep)
forall a. Rename a => a -> RenameM a
rename VName -> RenameM VName
forall a. Rename a => a -> RenameM a
rename

-- | The type of a SOAC.
soacType :: (Typed (LParamInfo rep)) => SOAC rep -> [ExtType]
soacType :: forall rep.
Typed (LParamInfo rep) =>
SOAC rep -> [TypeBase (ShapeBase ExtSize) NoMode]
soacType (JVP Shape
shape [SubExp]
_ [SubExp]
_ Lambda rep
lam) =
  [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall o. [TypeBase Shape o] -> [TypeBase (ShapeBase ExtSize) o]
staticShapes ([TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$
    Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam
      [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ (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` Shape
shape) (Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam)
soacType (VJP Shape
shape [SubExp]
_ [SubExp]
_ Lambda rep
lam) =
  [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall o. [TypeBase Shape o] -> [TypeBase (ShapeBase ExtSize) o]
staticShapes ([TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ (Param (LParamInfo rep) -> TypeBase Shape NoMode)
-> [Param (LParamInfo rep)] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
`arrayOfShape` Shape
shape) (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> (Param (LParamInfo rep) -> TypeBase Shape NoMode)
-> Param (LParamInfo 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
. Param (LParamInfo rep) -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
paramType) (Lambda rep -> [Param (LParamInfo rep)]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams Lambda rep
lam)
soacType (WithVJP [SubExp]
_ Lambda rep
lam Lambda rep
_) =
  [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall o. [TypeBase Shape o] -> [TypeBase (ShapeBase ExtSize) o]
staticShapes ([TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam
soacType (Stream SubExp
outersize [VName]
_ [SubExp]
accs Lambda rep
lam) =
  [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall o. [TypeBase Shape o] -> [TypeBase (ShapeBase ExtSize) o]
staticShapes ([TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$ (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map (Map VName SubExp -> TypeBase Shape NoMode -> TypeBase Shape NoMode
substNamesInType Map VName SubExp
substs) [TypeBase Shape NoMode]
rtp
  where
    nms :: [VName]
nms = (Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName ([Param (LParamInfo rep)] -> [VName])
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. Int -> [a] -> [a]
take (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
accs) [Param (LParamInfo rep)]
params
    substs :: Map VName SubExp
substs = [(VName, SubExp)] -> Map VName SubExp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, SubExp)] -> Map VName SubExp)
-> [(VName, SubExp)] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ [VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
nms (SubExp
outersize SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
accs)
    Lambda [Param (LParamInfo rep)]
params [TypeBase Shape NoMode]
rtp Body rep
_ = Lambda rep
lam
soacType (Hist SubExp
_ [VName]
_ [HistOp rep]
ops Lambda rep
_bucket_fun) = [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall o. [TypeBase Shape o] -> [TypeBase (ShapeBase ExtSize) o]
staticShapes ([TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$ do
  op <- [HistOp rep]
ops
  map (`arrayOfShape` histShape op) (lambdaReturnType $ histOp op)
soacType (Screma SubExp
w [VName]
_arrs ScremaForm rep
form) =
  [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall o. [TypeBase Shape o] -> [TypeBase (ShapeBase ExtSize) o]
staticShapes ([TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm rep -> [TypeBase Shape NoMode]
forall rep. SubExp -> ScremaForm rep -> [TypeBase Shape NoMode]
scremaType SubExp
w ScremaForm rep
form
soacType (FlatMap SubExp
w [VName]
_ ExtLambda rep
lam) =
  [ PrimType -> TypeBase (ShapeBase ExtSize) NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64,
    TypeBase (ShapeBase ExtSize) NoMode
-> ExtSize -> TypeBase (ShapeBase ExtSize) NoMode
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow (PrimType -> TypeBase (ShapeBase ExtSize) NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) (SubExp -> ExtSize
forall a. a -> Ext a
Free SubExp
w),
    TypeBase (ShapeBase ExtSize) NoMode
-> ExtSize -> TypeBase (ShapeBase ExtSize) NoMode
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow (PrimType -> TypeBase (ShapeBase ExtSize) NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool) (Int -> ExtSize
forall a. Int -> Ext a
Ext Int
0),
    TypeBase (ShapeBase ExtSize) NoMode
-> ExtSize -> TypeBase (ShapeBase ExtSize) NoMode
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow (PrimType -> TypeBase (ShapeBase ExtSize) NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) (SubExp -> ExtSize
forall a. a -> Ext a
Free SubExp
w)
  ]
    [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a. [a] -> [a] -> [a]
++ (TypeBase (ShapeBase ExtSize) NoMode
 -> TypeBase (ShapeBase ExtSize) NoMode)
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase (ShapeBase ExtSize) NoMode
-> TypeBase (ShapeBase ExtSize) NoMode
onValueRet (Int
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a. Int -> [a] -> [a]
drop Int
1 (ExtLambda rep -> [TypeBase (ShapeBase ExtSize) NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda rep
lam))
  where
    onValueRet :: TypeBase (ShapeBase ExtSize) NoMode
-> TypeBase (ShapeBase ExtSize) NoMode
onValueRet TypeBase (ShapeBase ExtSize) NoMode
t
      -- An irregular result is concatenated, so its outer size becomes the
      -- total size @m@.
      | TypeBase (ShapeBase ExtSize) NoMode -> Bool
flatMapNonuniform TypeBase (ShapeBase ExtSize) NoMode
t = TypeBase (ShapeBase ExtSize) NoMode
t TypeBase (ShapeBase ExtSize) NoMode
-> ExtSize -> TypeBase (ShapeBase ExtSize) NoMode
forall d m.
TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
`setOuterSize` Int -> ExtSize
forall a. Int -> Ext a
Ext Int
0
      -- A regular result is merely collected, one per input element.
      | Bool
otherwise = TypeBase (ShapeBase ExtSize) NoMode
t TypeBase (ShapeBase ExtSize) NoMode
-> ExtSize -> TypeBase (ShapeBase ExtSize) NoMode
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
`arrayOfRow` SubExp -> ExtSize
forall a. a -> Ext a
Free SubExp
w

-- | Is this the type of an irregular result of a 'FlatMap' lambda - one that is
-- concatenated with the results of the other iterations, rather than merely
-- collected? These are exactly the results whose outermost size is the
-- existential size returned by the lambda.
flatMapNonuniform :: ExtType -> Bool
flatMapNonuniform :: TypeBase (ShapeBase ExtSize) NoMode -> Bool
flatMapNonuniform TypeBase (ShapeBase ExtSize) NoMode
t =
  case ShapeBase ExtSize -> [ExtSize]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase ExtSize -> [ExtSize]) -> ShapeBase ExtSize -> [ExtSize]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase ExtSize) NoMode -> ShapeBase ExtSize
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase (ShapeBase ExtSize) NoMode
t of
    Ext Int
0 : [ExtSize]
_ -> Bool
True
    [ExtSize]
_ -> Bool
False

-- | The element types of the irregular (concatenated) arrays produced by a
-- 'FlatMap' lambda, with the existential outer dimension stripped. The
-- remaining dimensions are never existential.
flatMapRowTypes :: ExtLambda rep -> [Type]
flatMapRowTypes :: forall rep. ExtLambda rep -> [TypeBase Shape NoMode]
flatMapRowTypes ExtLambda rep
lam =
  [TypeBase Shape NoMode]
-> Maybe [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. a -> Maybe a -> a
fromMaybe (String -> [TypeBase Shape NoMode]
forall a. HasCallStack => String -> a
error String
"flatMapRowTypes: existential inner dimension.") (Maybe [TypeBase Shape NoMode] -> [TypeBase Shape NoMode])
-> Maybe [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$
    (TypeBase (ShapeBase ExtSize) NoMode
 -> Maybe (TypeBase Shape NoMode))
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> Maybe [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 (TypeBase (ShapeBase ExtSize) NoMode
-> Maybe (TypeBase Shape NoMode)
forall o.
TypeBase (ShapeBase ExtSize) o -> Maybe (TypeBase Shape o)
hasStaticShape (TypeBase (ShapeBase ExtSize) NoMode
 -> Maybe (TypeBase Shape NoMode))
-> (TypeBase (ShapeBase ExtSize) NoMode
    -> TypeBase (ShapeBase ExtSize) NoMode)
-> TypeBase (ShapeBase ExtSize) NoMode
-> Maybe (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
. TypeBase (ShapeBase ExtSize) NoMode
-> TypeBase (ShapeBase ExtSize) NoMode
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType) ([TypeBase (ShapeBase ExtSize) NoMode]
 -> Maybe [TypeBase Shape NoMode])
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> Maybe [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$
      (TypeBase (ShapeBase ExtSize) NoMode -> Bool)
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a. (a -> Bool) -> [a] -> [a]
filter TypeBase (ShapeBase ExtSize) NoMode -> Bool
flatMapNonuniform ([TypeBase (ShapeBase ExtSize) NoMode]
 -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$
        Int
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a. Int -> [a] -> [a]
drop Int
1 (ExtLambda rep -> [TypeBase (ShapeBase ExtSize) NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda rep
lam)

-- | Split a list with one element per result of a 'FlatMap' into its metadata
-- results - the size, the shape array, the flag array, and the offset array -
-- and its value results.
flatMapSplitMeta :: [a] -> ((a, a, a, a), [a])
flatMapSplitMeta :: forall a. [a] -> ((a, a, a, a), [a])
flatMapSplitMeta (a
m : a
shape : a
flags : a
offset : [a]
vals) =
  ((a
m, a
shape, a
flags, a
offset), [a]
vals)
flatMapSplitMeta [a]
_ =
  String -> ((a, a, a, a), [a])
forall a. HasCallStack => String -> a
error String
"flatMapSplitMeta: too few results."

-- | Split a list with one element per value result of a 'FlatMap' lambda into
-- the irregular and the regular parts.
flatMapSplitValues :: ExtLambda rep -> [a] -> ([a], [a])
flatMapSplitValues :: forall rep a. ExtLambda rep -> [a] -> ([a], [a])
flatMapSplitValues ExtLambda rep
lam =
  [Either a a] -> ([a], [a])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either a a] -> ([a], [a]))
-> ([a] -> [Either a a]) -> [a] -> ([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
. (TypeBase (ShapeBase ExtSize) NoMode -> a -> Either a a)
-> [TypeBase (ShapeBase ExtSize) NoMode] -> [a] -> [Either a a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase (ShapeBase ExtSize) NoMode -> a -> Either a a
forall {b}. TypeBase (ShapeBase ExtSize) NoMode -> b -> Either b b
f (Int
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a. Int -> [a] -> [a]
drop Int
1 (ExtLambda rep -> [TypeBase (ShapeBase ExtSize) NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda rep
lam))
  where
    f :: TypeBase (ShapeBase ExtSize) NoMode -> b -> Either b b
f TypeBase (ShapeBase ExtSize) NoMode
t b
x = if TypeBase (ShapeBase ExtSize) NoMode -> Bool
flatMapNonuniform TypeBase (ShapeBase ExtSize) NoMode
t then b -> Either b b
forall a b. a -> Either a b
Left b
x else b -> Either b b
forall a b. b -> Either a b
Right b
x

-- | The types of the regular results of a 'FlatMap' lambda - the ones that are
-- collected rather than concatenated. These never have existential sizes.
flatMapUniformTypes :: ExtLambda rep -> [Type]
flatMapUniformTypes :: forall rep. ExtLambda rep -> [TypeBase Shape NoMode]
flatMapUniformTypes ExtLambda rep
lam =
  [TypeBase Shape NoMode]
-> Maybe [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. a -> Maybe a -> a
fromMaybe (String -> [TypeBase Shape NoMode]
forall a. HasCallStack => String -> a
error String
"flatMapUniformTypes: existential size.") (Maybe [TypeBase Shape NoMode] -> [TypeBase Shape NoMode])
-> Maybe [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$
    (TypeBase (ShapeBase ExtSize) NoMode
 -> Maybe (TypeBase Shape NoMode))
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> Maybe [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 TypeBase (ShapeBase ExtSize) NoMode
-> Maybe (TypeBase Shape NoMode)
forall o.
TypeBase (ShapeBase ExtSize) o -> Maybe (TypeBase Shape o)
hasStaticShape ([TypeBase (ShapeBase ExtSize) NoMode]
 -> Maybe [TypeBase Shape NoMode])
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> Maybe [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$
      (TypeBase (ShapeBase ExtSize) NoMode -> Bool)
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> (TypeBase (ShapeBase ExtSize) NoMode -> Bool)
-> TypeBase (ShapeBase ExtSize) NoMode
-> 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
. TypeBase (ShapeBase ExtSize) NoMode -> Bool
flatMapNonuniform) ([TypeBase (ShapeBase ExtSize) NoMode]
 -> [TypeBase (ShapeBase ExtSize) NoMode])
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a b. (a -> b) -> a -> b
$
        Int
-> [TypeBase (ShapeBase ExtSize) NoMode]
-> [TypeBase (ShapeBase ExtSize) NoMode]
forall a. Int -> [a] -> [a]
drop Int
1 (ExtLambda rep -> [TypeBase (ShapeBase ExtSize) NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda rep
lam)

instance TypedOp SOAC where
  opType :: forall rep (m :: * -> *).
HasScope rep m =>
SOAC rep -> m [TypeBase (ShapeBase ExtSize) NoMode]
opType = [TypeBase (ShapeBase ExtSize) NoMode]
-> m [TypeBase (ShapeBase ExtSize) NoMode]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TypeBase (ShapeBase ExtSize) NoMode]
 -> m [TypeBase (ShapeBase ExtSize) NoMode])
-> (SOAC rep -> [TypeBase (ShapeBase ExtSize) NoMode])
-> SOAC rep
-> m [TypeBase (ShapeBase ExtSize) 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
. SOAC rep -> [TypeBase (ShapeBase ExtSize) NoMode]
forall rep.
Typed (LParamInfo rep) =>
SOAC rep -> [TypeBase (ShapeBase ExtSize) NoMode]
soacType

instance AliasedOp SOAC where
  opAliases :: forall rep. Aliased rep => SOAC rep -> [Names]
opAliases = (TypeBase (ShapeBase ExtSize) NoMode -> Names)
-> [TypeBase (ShapeBase ExtSize) NoMode] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (Names -> TypeBase (ShapeBase ExtSize) NoMode -> Names
forall a b. a -> b -> a
const Names
forall a. Monoid a => a
mempty) ([TypeBase (ShapeBase ExtSize) NoMode] -> [Names])
-> (SOAC rep -> [TypeBase (ShapeBase ExtSize) NoMode])
-> SOAC 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
. SOAC rep -> [TypeBase (ShapeBase ExtSize) NoMode]
forall rep.
Typed (LParamInfo rep) =>
SOAC rep -> [TypeBase (ShapeBase ExtSize) NoMode]
soacType

  consumedInOp :: forall rep. Aliased rep => SOAC rep -> Names
consumedInOp JVP {} = Names
forall a. Monoid a => a
mempty
  consumedInOp VJP {} = Names
forall a. Monoid a => a
mempty
  consumedInOp WithVJP {} = Names
forall a. Monoid a => a
mempty
  consumedInOp FlatMap {} = Names
forall a. Monoid a => a
mempty
  -- Only map functions can consume anything.  The operands to scan
  -- and reduce functions are always considered "fresh".
  consumedInOp (Screma SubExp
_ [VName]
arrs (ScremaForm Lambda rep
map_lam [Scan rep]
_ [Reduce rep]
_ Lambda rep
_)) =
    (VName -> VName) -> Names -> Names
mapNames VName -> VName
consumedArray (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Names
forall rep. Aliased rep => Lambda rep -> Names
consumedByLambda Lambda rep
map_lam
    where
      consumedArray :: VName -> VName
consumedArray VName
v = VName -> Maybe VName -> VName
forall a. a -> Maybe a -> a
fromMaybe VName
v (Maybe VName -> VName) -> Maybe VName -> VName
forall a b. (a -> b) -> a -> b
$ VName -> [(VName, VName)] -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v [(VName, VName)]
params_to_arrs
      params_to_arrs :: [(VName, VName)]
params_to_arrs = [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName ([Param (LParamInfo rep)] -> [VName])
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Param (LParamInfo rep)]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams Lambda rep
map_lam) [VName]
arrs
  consumedInOp (Stream SubExp
_ [VName]
arrs [SubExp]
accs Lambda rep
lam) =
    [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [VName]
subExpVars ([SubExp] -> [VName]) -> [SubExp] -> [VName]
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
consumedArray ([VName] -> [SubExp]) -> [VName] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Names
forall rep. Aliased rep => Lambda rep -> Names
consumedByLambda Lambda rep
lam
    where
      consumedArray :: VName -> SubExp
consumedArray VName
v = SubExp -> Maybe SubExp -> SubExp
forall a. a -> Maybe a -> a
fromMaybe (VName -> SubExp
Var VName
v) (Maybe SubExp -> SubExp) -> Maybe SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ VName -> [(VName, SubExp)] -> Maybe SubExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v [(VName, SubExp)]
paramsToInput
      -- Drop the chunk parameter, which cannot alias anything.
      paramsToInput :: [(VName, SubExp)]
paramsToInput =
        [VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName ([Param (LParamInfo rep)] -> [VName])
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. Int -> [a] -> [a]
drop Int
1 ([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 -> [Param (LParamInfo rep)]
lambdaParams Lambda rep
lam) ([SubExp]
accs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
arrs)
  consumedInOp (Hist SubExp
_ [VName]
_ [HistOp rep]
ops Lambda rep
_) =
    [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (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

mapHistOp ::
  (Lambda frep -> Lambda trep) ->
  HistOp frep ->
  HistOp trep
mapHistOp :: forall frep trep.
(Lambda frep -> Lambda trep) -> HistOp frep -> HistOp trep
mapHistOp Lambda frep -> Lambda trep
f (HistOp Shape
w SubExp
rf [VName]
dests [SubExp]
nes Lambda frep
lam) =
  Shape
-> SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep
forall rep.
Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
HistOp Shape
w SubExp
rf [VName]
dests [SubExp]
nes (Lambda trep -> HistOp trep) -> Lambda trep -> HistOp trep
forall a b. (a -> b) -> a -> b
$ Lambda frep -> Lambda trep
f Lambda frep
lam

instance CanBeAliased SOAC where
  addOpAliases :: forall rep.
AliasableRep rep =>
AliasTable -> SOAC rep -> SOAC (Aliases rep)
addOpAliases AliasTable
aliases (JVP Shape
shape [SubExp]
args [SubExp]
vec Lambda rep
lam) =
    Shape
-> [SubExp]
-> [SubExp]
-> Lambda (Aliases rep)
-> SOAC (Aliases rep)
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
JVP Shape
shape [SubExp]
args [SubExp]
vec (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
lam)
  addOpAliases AliasTable
aliases (VJP Shape
shape [SubExp]
args [SubExp]
vec Lambda rep
lam) =
    Shape
-> [SubExp]
-> [SubExp]
-> Lambda (Aliases rep)
-> SOAC (Aliases rep)
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
VJP Shape
shape [SubExp]
args [SubExp]
vec (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
lam)
  addOpAliases AliasTable
aliases (WithVJP [SubExp]
args Lambda rep
lam Lambda rep
lam_adj) =
    [SubExp]
-> Lambda (Aliases rep)
-> Lambda (Aliases rep)
-> SOAC (Aliases rep)
forall rep. [SubExp] -> Lambda rep -> Lambda rep -> SOAC rep
WithVJP
      [SubExp]
args
      (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
lam)
      (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
lam_adj)
  addOpAliases AliasTable
aliases (FlatMap SubExp
w [VName]
arrs ExtLambda rep
lam) =
    SubExp -> [VName] -> ExtLambda (Aliases rep) -> SOAC (Aliases rep)
forall rep. SubExp -> [VName] -> ExtLambda rep -> SOAC rep
FlatMap SubExp
w [VName]
arrs (ExtLambda (Aliases rep) -> SOAC (Aliases rep))
-> ExtLambda (Aliases rep) -> SOAC (Aliases rep)
forall a b. (a -> b) -> a -> b
$ AliasTable -> ExtLambda rep -> ExtLambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases ExtLambda rep
lam
  addOpAliases AliasTable
aliases (Stream SubExp
size [VName]
arr [SubExp]
accs Lambda rep
lam) =
    SubExp
-> [VName]
-> [SubExp]
-> Lambda (Aliases rep)
-> SOAC (Aliases rep)
forall rep. SubExp -> [VName] -> [SubExp] -> Lambda rep -> SOAC rep
Stream SubExp
size [VName]
arr [SubExp]
accs (Lambda (Aliases rep) -> SOAC (Aliases rep))
-> Lambda (Aliases rep) -> SOAC (Aliases rep)
forall a b. (a -> b) -> a -> b
$ AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
lam
  addOpAliases AliasTable
aliases (Hist SubExp
w [VName]
arrs [HistOp rep]
ops Lambda rep
bucket_fun) =
    SubExp
-> [VName]
-> [HistOp (Aliases rep)]
-> Lambda (Aliases rep)
-> SOAC (Aliases rep)
forall rep.
SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
Hist
      SubExp
w
      [VName]
arrs
      ((HistOp rep -> HistOp (Aliases rep))
-> [HistOp rep] -> [HistOp (Aliases rep)]
forall a b. (a -> b) -> [a] -> [b]
map ((Lambda rep -> Lambda (Aliases rep))
-> HistOp rep -> HistOp (Aliases rep)
forall frep trep.
(Lambda frep -> Lambda trep) -> HistOp frep -> HistOp trep
mapHistOp (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases)) [HistOp rep]
ops)
      (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
bucket_fun)
  addOpAliases AliasTable
aliases (Screma SubExp
w [VName]
arrs (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam)) =
    SubExp -> [VName] -> ScremaForm (Aliases rep) -> SOAC (Aliases rep)
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs (ScremaForm (Aliases rep) -> SOAC (Aliases rep))
-> ScremaForm (Aliases rep) -> SOAC (Aliases rep)
forall a b. (a -> b) -> a -> b
$
      Lambda (Aliases rep)
-> [Scan (Aliases rep)]
-> [Reduce (Aliases rep)]
-> Lambda (Aliases rep)
-> ScremaForm (Aliases rep)
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm
        (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
map_lam)
        ((Scan rep -> Scan (Aliases rep))
-> [Scan rep] -> [Scan (Aliases rep)]
forall a b. (a -> b) -> [a] -> [b]
map Scan rep -> Scan (Aliases rep)
onScan [Scan rep]
scans)
        ((Reduce rep -> Reduce (Aliases rep))
-> [Reduce rep] -> [Reduce (Aliases rep)]
forall a b. (a -> b) -> [a] -> [b]
map Reduce rep -> Reduce (Aliases rep)
onRed [Reduce rep]
reds)
        (AliasTable -> Lambda rep -> Lambda (Aliases rep)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
aliases Lambda rep
post_lam)
    where
      onRed :: Reduce rep -> Reduce (Aliases rep)
onRed Reduce rep
red = Reduce rep
red {redLambda = Alias.analyseLambda aliases $ redLambda red}
      onScan :: Scan rep -> Scan (Aliases rep)
onScan Scan rep
scan = Scan rep
scan {scanLambda = Alias.analyseLambda aliases $ scanLambda scan}

instance IsOp SOAC where
  safeOp :: forall rep. ASTRep rep => SOAC rep -> Bool
safeOp SOAC rep
_ = Bool
False
  cheapOp :: forall rep. ASTRep rep => SOAC rep -> Bool
cheapOp SOAC rep
_ = Bool
False
  opDependencies :: forall rep. ASTRep rep => SOAC rep -> [Names]
opDependencies (Stream SubExp
w [VName]
arrs [SubExp]
accs Lambda rep
lam) =
    let accs_deps :: [Names]
accs_deps = (SubExp -> Names) -> [SubExp] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Names
depsOf' [SubExp]
accs
        arrs_deps :: [Names]
arrs_deps = SubExp -> [VName] -> [Names]
depsOfArrays SubExp
w [VName]
arrs
     in AliasTable -> Lambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
lam ([Names]
arrs_deps [Names] -> [Names] -> [Names]
forall a. Semigroup a => a -> a -> a
<> [Names]
accs_deps)
  opDependencies (Hist SubExp
w [VName]
arrs [HistOp rep]
ops Lambda rep
lam) =
    let bucket_fun_deps' :: [Names]
bucket_fun_deps' = AliasTable -> Lambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
lam (SubExp -> [VName] -> [Names]
depsOfArrays SubExp
w [VName]
arrs)
        -- Bucket function results are indices followed by values.
        -- Reshape this to align with list of histogram operations.
        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
        value_lengths :: [Int]
value_lengths = (HistOp rep -> Int) -> [HistOp 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) -> (HistOp rep -> [SubExp]) -> 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 -> [SubExp]
forall rep. HistOp rep -> [SubExp]
histNeutral) [HistOp rep]
ops
        ([Names]
indices, [Names]
values) = Int -> [Names] -> ([Names], [Names])
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) [Names]
bucket_fun_deps'
        bucket_fun_deps :: [[Names]]
bucket_fun_deps =
          ([Names] -> [Names] -> [Names])
-> [[Names]] -> [[Names]] -> [[Names]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
            [Names] -> [Names] -> [Names]
forall {b}. Monoid b => [b] -> [b] -> [b]
concatIndicesToEachValue
            ([Int] -> [Names] -> [[Names]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
ranks [Names]
indices)
            ([Int] -> [Names] -> [[Names]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
value_lengths [Names]
values)
     in [[Names]] -> [Names]
forall a. Monoid a => [a] -> a
mconcat ([[Names]] -> [Names]) -> [[Names]] -> [Names]
forall a b. (a -> b) -> a -> b
$ ([Names] -> [Names] -> [Names])
-> [[Names]] -> [[Names]] -> [[Names]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((Names -> Names -> Names) -> [Names] -> [Names] -> [Names]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
(<>)) [[Names]]
bucket_fun_deps ((HistOp rep -> [Names]) -> [HistOp rep] -> [[Names]]
forall a b. (a -> b) -> [a] -> [b]
map HistOp rep -> [Names]
forall {rep}. ASTRep rep => HistOp rep -> [Names]
depsOfHistOp [HistOp rep]
ops)
    where
      depsOfHistOp :: HistOp rep -> [Names]
depsOfHistOp (HistOp Shape
dest_shape SubExp
rf [VName]
dests [SubExp]
nes Lambda rep
op) =
        let shape_deps :: Names
shape_deps = Shape -> Names
depsOfShape Shape
dest_shape
            in_deps :: [Names]
in_deps = (VName -> Names) -> [VName] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (\VName
vn -> VName -> Names
oneName VName
vn Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
shape_deps Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> SubExp -> Names
depsOf' SubExp
rf) [VName]
dests
         in AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
forall rep.
ASTRep rep =>
AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
reductionDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
op [SubExp]
nes [Names]
in_deps
      -- A histogram operation may use the same index for multiple values.
      concatIndicesToEachValue :: [b] -> [b] -> [b]
concatIndicesToEachValue [b]
is [b]
vs =
        let is_flat :: b
is_flat = [b] -> b
forall a. Monoid a => [a] -> a
mconcat [b]
is
         in (b -> b) -> [b] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (b
is_flat b -> b -> b
forall a. Semigroup a => a -> a -> a
<>) [b]
vs
  opDependencies (JVP Shape
_ [SubExp]
args [SubExp]
vec Lambda rep
lam) =
    [[Names]] -> [Names]
forall a. Monoid a => [a] -> a
mconcat ([[Names]] -> [Names]) -> [[Names]] -> [Names]
forall a b. (a -> b) -> a -> b
$
      Int -> [Names] -> [[Names]]
forall a. Int -> a -> [a]
replicate Int
2 ([Names] -> [[Names]]) -> [Names] -> [[Names]]
forall a b. (a -> b) -> a -> b
$
        AliasTable -> Lambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
lam ([Names] -> [Names]) -> [Names] -> [Names]
forall a b. (a -> b) -> a -> b
$
          (Names -> Names -> Names) -> [Names] -> [Names] -> [Names]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
(<>) ((SubExp -> Names) -> [SubExp] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Names
depsOf' [SubExp]
args) ((SubExp -> Names) -> [SubExp] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Names
depsOf' [SubExp]
vec)
  opDependencies (VJP Shape
_ [SubExp]
args [SubExp]
vec Lambda rep
lam) =
    AliasTable -> Lambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies
      AliasTable
forall a. Monoid a => a
mempty
      Lambda rep
lam
      ((Names -> Names -> Names) -> [Names] -> [Names] -> [Names]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
(<>) ((SubExp -> Names) -> [SubExp] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Names
depsOf' [SubExp]
args) ((SubExp -> Names) -> [SubExp] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Names
depsOf' [SubExp]
vec))
      [Names] -> [Names] -> [Names]
forall a. Semigroup a => a -> a -> a
<> (Param (LParamInfo rep) -> Names)
-> [Param (LParamInfo rep)] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (Names -> Param (LParamInfo rep) -> Names
forall a b. a -> b -> a
const (Names -> Param (LParamInfo rep) -> Names)
-> Names -> Param (LParamInfo rep) -> Names
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
args Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda rep
lam) (Lambda rep -> [Param (LParamInfo rep)]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams Lambda rep
lam)
  opDependencies (WithVJP [SubExp]
args Lambda rep
lam Lambda rep
_lam_adj) =
    AliasTable -> Lambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies
      AliasTable
forall a. Monoid a => a
mempty
      Lambda rep
lam
      ((SubExp -> Names) -> [SubExp] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Names
depsOf' [SubExp]
args)
      [Names] -> [Names] -> [Names]
forall a. Semigroup a => a -> a -> a
<> (Param (LParamInfo rep) -> Names)
-> [Param (LParamInfo rep)] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (Names -> Param (LParamInfo rep) -> Names
forall a b. a -> b -> a
const (Names -> Param (LParamInfo rep) -> Names)
-> Names -> Param (LParamInfo rep) -> Names
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
args Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda rep
lam) (Lambda rep -> [Param (LParamInfo rep)]
forall {rep} {t}. GLambda rep t -> [Param (LParamInfo rep)]
lambdaParams Lambda rep
lam)
  opDependencies (FlatMap SubExp
w [VName]
arrs ExtLambda rep
lam) =
    -- FIXME: this is strictly speaking not accurate, although it will not cause
    -- trouble anytime soon.
    Int -> Names -> [Names]
forall a. Int -> a -> [a]
replicate Int
4 Names
forall a. Monoid a => a
mempty
      [Names] -> [Names] -> [Names]
forall a. [a] -> [a] -> [a]
++ Int -> [Names] -> [Names]
forall a. Int -> [a] -> [a]
drop Int
1 (AliasTable -> ExtLambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty ExtLambda rep
lam (SubExp -> [VName] -> [Names]
depsOfArrays SubExp
w [VName]
arrs))
  opDependencies (Screma SubExp
w [VName]
arrs (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam)) =
    let ([Names]
scans_in, [Names]
reds_in, [Names]
map_deps) =
          Int -> Int -> [Names] -> ([Names], [Names], [Names])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 ([Scan rep] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan rep]
scans) ([Reduce rep] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce rep]
reds) ([Names] -> ([Names], [Names], [Names]))
-> [Names] -> ([Names], [Names], [Names])
forall a b. (a -> b) -> a -> b
$
            AliasTable -> Lambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
map_lam (SubExp -> [VName] -> [Names]
depsOfArrays SubExp
w [VName]
arrs)
        scans_deps :: [Names]
scans_deps =
          ((Scan rep, [Names]) -> [Names])
-> [(Scan rep, [Names])] -> [Names]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Scan rep, [Names]) -> [Names]
forall {rep}. ASTRep rep => (Scan rep, [Names]) -> [Names]
depsOfScan ([Scan rep] -> [[Names]] -> [(Scan rep, [Names])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Scan rep]
scans ([[Names]] -> [(Scan rep, [Names])])
-> [[Names]] -> [(Scan rep, [Names])]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Names] -> [[Names]]
forall a. [Int] -> [a] -> [[a]]
chunks ([Scan rep] -> [Int]
forall rep. [Scan rep] -> [Int]
scanSizes [Scan rep]
scans) [Names]
scans_in)
        reds_deps :: [Names]
reds_deps =
          ((Reduce rep, [Names]) -> [Names])
-> [(Reduce rep, [Names])] -> [Names]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Reduce rep, [Names]) -> [Names]
forall {rep}. ASTRep rep => (Reduce rep, [Names]) -> [Names]
depsOfRed ([Reduce rep] -> [[Names]] -> [(Reduce rep, [Names])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Reduce rep]
reds ([[Names]] -> [(Reduce rep, [Names])])
-> [[Names]] -> [(Reduce rep, [Names])]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Names] -> [[Names]]
forall a. [Int] -> [a] -> [[a]]
chunks ([Reduce rep] -> [Int]
forall rep. [Reduce rep] -> [Int]
redSizes [Reduce rep]
reds) [Names]
reds_in)
     in [Names]
reds_deps [Names] -> [Names] -> [Names]
forall a. Semigroup a => a -> a -> a
<> AliasTable -> Lambda rep -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
post_lam ([Names]
scans_deps [Names] -> [Names] -> [Names]
forall a. Semigroup a => a -> a -> a
<> [Names]
map_deps)
    where
      depsOfScan :: (Scan rep, [Names]) -> [Names]
depsOfScan (Scan Lambda rep
lam [SubExp]
nes, [Names]
deps_in) =
        AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
forall rep.
ASTRep rep =>
AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
reductionDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
lam [SubExp]
nes [Names]
deps_in
      depsOfRed :: (Reduce rep, [Names]) -> [Names]
depsOfRed (Reduce Commutativity
_ Lambda rep
lam [SubExp]
nes, [Names]
deps_in) =
        AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
forall rep.
ASTRep rep =>
AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
reductionDependencies AliasTable
forall a. Monoid a => a
mempty Lambda rep
lam [SubExp]
nes [Names]
deps_in

substNamesInType :: M.Map VName SubExp -> Type -> Type
substNamesInType :: Map VName SubExp -> TypeBase Shape NoMode -> TypeBase Shape NoMode
substNamesInType Map VName SubExp
_ t :: TypeBase Shape NoMode
t@Prim {} = TypeBase Shape NoMode
t
substNamesInType Map VName SubExp
_ t :: TypeBase Shape NoMode
t@Acc {} = TypeBase Shape NoMode
t
substNamesInType Map VName SubExp
_ (Mem Space
space) = Space -> TypeBase Shape NoMode
forall shape o. Space -> TypeBase shape o
Mem Space
space
substNamesInType Map VName SubExp
subs (Array PrimType
btp Shape
shp NoMode
o) =
  let shp' :: Shape
shp' = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ (SubExp -> SubExp) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (Map VName SubExp -> SubExp -> SubExp
substNamesInSubExp Map VName SubExp
subs) (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shp)
   in PrimType -> Shape -> NoMode -> TypeBase Shape NoMode
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
btp Shape
shp' NoMode
o

substNamesInSubExp :: M.Map VName SubExp -> SubExp -> SubExp
substNamesInSubExp :: Map VName SubExp -> SubExp -> SubExp
substNamesInSubExp Map VName SubExp
_ e :: SubExp
e@(Constant PrimValue
_) = SubExp
e
substNamesInSubExp Map VName SubExp
subs (Var VName
idd) =
  SubExp -> VName -> Map VName SubExp -> SubExp
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault (VName -> SubExp
Var VName
idd) VName
idd Map VName SubExp
subs

instance CanBeWise SOAC where
  addOpWisdom :: forall rep. Informing rep => SOAC rep -> SOAC (Wise rep)
addOpWisdom = Identity (SOAC (Wise rep)) -> SOAC (Wise rep)
forall a. Identity a -> a
runIdentity (Identity (SOAC (Wise rep)) -> SOAC (Wise rep))
-> (SOAC rep -> Identity (SOAC (Wise rep)))
-> SOAC rep
-> SOAC (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
. SOACMapper rep (Wise rep) Identity
-> SOAC rep -> Identity (SOAC (Wise rep))
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM ((SubExp -> Identity SubExp)
-> (Lambda rep -> Identity (Lambda (Wise rep)))
-> (ExtLambda rep -> Identity (ExtLambda (Wise rep)))
-> (VName -> Identity VName)
-> SOACMapper rep (Wise rep) Identity
forall frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Lambda frep -> m (Lambda trep))
-> (ExtLambda frep -> m (ExtLambda trep))
-> (VName -> m VName)
-> SOACMapper frep trep m
SOACMapper 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) (ExtLambda (Wise rep) -> Identity (ExtLambda (Wise rep))
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExtLambda (Wise rep) -> Identity (ExtLambda (Wise rep)))
-> (ExtLambda rep -> ExtLambda (Wise rep))
-> ExtLambda rep
-> Identity (ExtLambda (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
. ExtLambda rep -> ExtLambda (Wise rep)
forall rep t.
Informing rep =>
GLambda rep t -> GLambda (Wise rep) t
informLambda) VName -> Identity VName
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure)

instance (RepTypes rep) => ST.IndexOp (SOAC rep) where
  indexOp :: forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int -> SOAC rep -> [TPrimExp Int64 VName] -> Maybe Indexed
indexOp SymbolTable rep
vtable Int
k SOAC rep
soac [TPrimExp Int64 VName
i] = do
    (lam, se, arr_params, arrs) <- SOAC rep
-> Maybe
     (GLambda rep (TypeBase Shape NoMode), SubExpRes,
      [Param (LParamInfo rep)], [VName])
lambdaAndSubExp SOAC rep
soac
    let arr_indexes = [(VName, (PrimExp VName, Certs))]
-> Map VName (PrimExp VName, Certs)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, (PrimExp VName, Certs))]
 -> Map VName (PrimExp VName, Certs))
-> [(VName, (PrimExp VName, Certs))]
-> Map VName (PrimExp VName, Certs)
forall a b. (a -> b) -> a -> b
$ [Maybe (VName, (PrimExp VName, Certs))]
-> [(VName, (PrimExp VName, Certs))]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (VName, (PrimExp VName, Certs))]
 -> [(VName, (PrimExp VName, Certs))])
-> [Maybe (VName, (PrimExp VName, Certs))]
-> [(VName, (PrimExp VName, Certs))]
forall a b. (a -> b) -> a -> b
$ (Param (LParamInfo rep)
 -> VName -> Maybe (VName, (PrimExp VName, Certs)))
-> [Param (LParamInfo rep)]
-> [VName]
-> [Maybe (VName, (PrimExp VName, Certs))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Param (LParamInfo rep)
-> VName -> Maybe (VName, (PrimExp VName, Certs))
arrIndex [Param (LParamInfo rep)]
arr_params [VName]
arrs
        arr_indexes' = (Map VName (PrimExp VName, Certs)
 -> Stm rep -> Map VName (PrimExp VName, Certs))
-> Map VName (PrimExp VName, Certs)
-> Seq (Stm rep)
-> Map VName (PrimExp VName, Certs)
forall b a. (b -> a -> b) -> b -> Seq a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map VName (PrimExp VName, Certs)
-> Stm rep -> Map VName (PrimExp VName, Certs)
expandPrimExpTable Map VName (PrimExp VName, Certs)
arr_indexes (Seq (Stm rep) -> Map VName (PrimExp VName, Certs))
-> Seq (Stm rep) -> Map VName (PrimExp VName, Certs)
forall a b. (a -> b) -> a -> b
$ GBody rep SubExpRes -> Seq (Stm rep)
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody rep SubExpRes -> Seq (Stm rep))
-> GBody rep SubExpRes -> Seq (Stm rep)
forall a b. (a -> b) -> a -> b
$ GLambda rep (TypeBase Shape NoMode) -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda rep (TypeBase Shape NoMode)
lam
    case se of
      SubExpRes Certs
_ (Var VName
v) -> (PrimExp VName -> Certs -> Indexed)
-> (PrimExp VName, Certs) -> Indexed
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Certs -> PrimExp VName -> Indexed)
-> PrimExp VName -> Certs -> Indexed
forall a b c. (a -> b -> c) -> b -> a -> c
flip Certs -> PrimExp VName -> Indexed
ST.Indexed) ((PrimExp VName, Certs) -> Indexed)
-> Maybe (PrimExp VName, Certs) -> Maybe Indexed
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName
-> Map VName (PrimExp VName, Certs) -> Maybe (PrimExp VName, Certs)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName (PrimExp VName, Certs)
arr_indexes'
      SubExpRes
_ -> Maybe Indexed
forall a. Maybe a
Nothing
    where
      lambdaAndSubExp :: SOAC rep
-> Maybe
     (GLambda rep (TypeBase Shape NoMode), SubExpRes,
      [Param (LParamInfo rep)], [VName])
lambdaAndSubExp (Screma SubExp
_ [VName]
arrs (ScremaForm GLambda rep (TypeBase Shape NoMode)
map_lam [Scan rep]
scans [Reduce rep]
reds GLambda rep (TypeBase Shape NoMode)
post_lam)) = do
        -- UNSURE_IF_CORRECT
        Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ GLambda rep (TypeBase Shape NoMode) -> Bool
forall rep. Lambda rep -> Bool
isIdentityLambda GLambda rep (TypeBase Shape NoMode)
post_lam
        Int
-> GLambda rep (TypeBase Shape NoMode)
-> [VName]
-> Maybe
     (GLambda rep (TypeBase Shape NoMode), SubExpRes,
      [Param (LParamInfo rep)], [VName])
nthMapOut ([Scan rep] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan rep]
scans Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Reduce rep] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce rep]
reds) GLambda rep (TypeBase Shape NoMode)
map_lam [VName]
arrs
      lambdaAndSubExp SOAC rep
_ =
        Maybe
  (GLambda rep (TypeBase Shape NoMode), SubExpRes,
   [Param (LParamInfo rep)], [VName])
forall a. Maybe a
Nothing

      nthMapOut :: Int
-> GLambda rep (TypeBase Shape NoMode)
-> [VName]
-> Maybe
     (GLambda rep (TypeBase Shape NoMode), SubExpRes,
      [Param (LParamInfo rep)], [VName])
nthMapOut Int
num_accs GLambda rep (TypeBase Shape NoMode)
lam [VName]
arrs = do
        se <- Int -> [SubExpRes] -> Maybe SubExpRes
forall int a. Integral int => int -> [a] -> Maybe a
maybeNth (Int
num_accs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) ([SubExpRes] -> Maybe SubExpRes) -> [SubExpRes] -> Maybe SubExpRes
forall a b. (a -> b) -> a -> b
$ GBody rep SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody rep SubExpRes -> [SubExpRes])
-> GBody rep SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ GLambda rep (TypeBase Shape NoMode) -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda rep (TypeBase Shape NoMode)
lam
        pure (lam, se, drop num_accs $ lambdaParams lam, arrs)

      arrIndex :: Param (LParamInfo rep)
-> VName -> Maybe (VName, (PrimExp VName, Certs))
arrIndex Param (LParamInfo rep)
p VName
arr = do
        ST.Indexed cs pe <- VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
ST.index' VName
arr [TPrimExp Int64 VName
i] SymbolTable rep
vtable
        pure (paramName p, (pe, cs))

      expandPrimExpTable :: Map VName (PrimExp VName, Certs)
-> Stm rep -> Map VName (PrimExp VName, Certs)
expandPrimExpTable Map VName (PrimExp VName, Certs)
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 (PrimExp VName, Certs)
-> VName -> WriterT Certs Maybe (PrimExp VName)
asPrimExp Map VName (PrimExp VName, Certs)
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 -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> SymbolTable rep -> Bool
forall rep. VName -> SymbolTable rep -> Bool
`ST.elem` SymbolTable rep
vtable) (Certs -> [VName]
unCerts (Certs -> [VName]) -> Certs -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm rep -> Certs
forall rep. Stm rep -> Certs
stmCerts Stm rep
stm) =
            VName
-> (PrimExp VName, Certs)
-> Map VName (PrimExp VName, Certs)
-> Map VName (PrimExp VName, Certs)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (PrimExp VName
pe, Stm rep -> Certs
forall rep. Stm rep -> Certs
stmCerts Stm rep
stm Certs -> Certs -> Certs
forall a. Semigroup a => a -> a -> a
<> Certs
cs) Map VName (PrimExp VName, Certs)
table
        | Bool
otherwise =
            Map VName (PrimExp VName, Certs)
table

      asPrimExp :: Map VName (PrimExp VName, Certs)
-> VName -> WriterT Certs Maybe (PrimExp VName)
asPrimExp Map VName (PrimExp VName, Certs)
table VName
v
        | Just (PrimExp VName
e, Certs
cs) <- VName
-> Map VName (PrimExp VName, Certs) -> Maybe (PrimExp VName, Certs)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName (PrimExp VName, Certs)
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
_ SOAC rep
_ [TPrimExp Int64 VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

-- | Type-check a SOAC.
typeCheckSOAC :: (TC.Checkable rep) => SOAC (Aliases rep) -> TC.TypeM rep ()
typeCheckSOAC :: forall rep. Checkable rep => SOAC (Aliases rep) -> TypeM rep ()
typeCheckSOAC (VJP Shape
shape [SubExp]
args [SubExp]
vec Lambda (Aliases rep)
lam) = 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 (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) Shape
shape
  args' <- (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]
args
  TC.checkLambda lam $ map TC.noArgAliases args'
  vec_ts <- mapM TC.checkSubExp vec
  unless (vec_ts == map (`arrayOfShape` shape) (lambdaReturnType lam)) $
    TC.bad . TC.TypeError . docText $
      "Return type"
        </> PP.indent 2 (pretty (lambdaReturnType lam))
        </> "inconsistent with type of seed vector"
        </> PP.indent 2 (pretty vec_ts)
        </> "with shape"
        </> PP.indent 2 (pretty shape)
typeCheckSOAC (JVP Shape
shape [SubExp]
args [SubExp]
vec Lambda (Aliases rep)
lam) = 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 (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) Shape
shape
  args' <- (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]
args
  TC.checkLambda lam $ map TC.noArgAliases args'
  vec_ts <- mapM TC.checkSubExp vec
  unless (vec_ts == map ((`arrayOfShape` shape) . TC.argType) args') $
    TC.bad . TC.TypeError . docText $
      "Parameter type"
        </> PP.indent 2 (pretty $ map TC.argType args')
        </> "does not match type of seed vector"
        </> PP.indent 2 (pretty vec_ts)
        </> "with shape"
        </> PP.indent 2 (pretty shape)
typeCheckSOAC (WithVJP [SubExp]
args Lambda (Aliases rep)
lam Lambda (Aliases rep)
lam_adj) = do
  args' <- (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]
args
  TC.checkLambda lam $ map TC.noArgAliases args'
  TC.checkLambda lam_adj $
    map (,mempty) (lambdaReturnType lam <> lambdaReturnType lam)
  unless (lambdaReturnType lam_adj == map TC.argType args') $
    TC.bad . TC.TypeError . docText $
      "Adjoint lambda return type"
        </> PP.indent 2 (pretty $ lambdaReturnType lam_adj)
        </> "does not match type of arguments"
        </> PP.indent 2 (pretty $ map TC.argType args')
typeCheckSOAC (FlatMap SubExp
w [VName]
arrs ExtLambda (Aliases rep)
lam) = 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
w
  arrs' <- SubExp -> [VName] -> TypeM rep [Arg]
forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
TC.checkSOACArrayArgs SubExp
w [VName]
arrs
  TC.checkExtLambda lam arrs'
typeCheckSOAC (Stream SubExp
size [VName]
arrexps [SubExp]
accexps Lambda (Aliases rep)
lam) = 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
size
  accargs <- (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]
accexps
  arrargs <- mapM lookupType arrexps
  _ <- TC.checkSOACArrayArgs size arrexps
  chunk <- case lambdaParams lam of
    LParam (Aliases rep)
chunk : [LParam (Aliases rep)]
_ -> Param (LParamInfo rep) -> TypeM rep (Param (LParamInfo rep))
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Param (LParamInfo rep)
LParam (Aliases rep)
chunk
    [] -> ErrorCase rep -> TypeM rep (Param (LParamInfo rep))
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep (Param (LParamInfo rep)))
-> ErrorCase rep -> TypeM rep (Param (LParamInfo rep))
forall a b. (a -> b) -> a -> b
$ Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TC.TypeError Text
"Stream lambda without parameters."
  let asArg a
t = (a
t, b
forall a. Monoid a => a
mempty)
      lamarrs' = (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Shape NoMode -> SubExp -> TypeBase Shape NoMode
forall d m.
TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
`setOuterSize` VName -> SubExp
Var (Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
chunk)) [TypeBase Shape NoMode]
arrargs
      acc_len = [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
accexps
      lamrtp = Int -> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Int -> [a] -> [a]
take Int
acc_len ([TypeBase Shape NoMode] -> [TypeBase Shape NoMode])
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$ Lambda (Aliases rep) -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda (Aliases rep)
lam
  unless (map TC.argType accargs == lamrtp) $
    TC.bad . TC.TypeError $
      "Stream with inconsistent accumulator type in lambda."
  -- just get the dflow of lambda on the fakearg, which does not alias
  -- arr, so we can later check that aliases of arr are not used inside lam.
  let fake_lamarrs' = (TypeBase Shape NoMode -> Arg) -> [TypeBase Shape NoMode] -> [Arg]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> Arg
forall {b} {a}. Monoid b => a -> (a, b)
asArg [TypeBase Shape NoMode]
lamarrs'
  TC.checkLambda lam $ asArg (Prim int64) : accargs ++ fake_lamarrs'
typeCheckSOAC (Hist SubExp
w [VName]
arrs [HistOp (Aliases rep)]
ops Lambda (Aliases rep)
bucket_fun) = 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
w

  -- Check the operators.
  [HistOp (Aliases rep)]
-> (HistOp (Aliases rep) -> TypeM rep ()) -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [HistOp (Aliases rep)]
ops ((HistOp (Aliases rep) -> TypeM rep ()) -> TypeM rep ())
-> (HistOp (Aliases rep) -> TypeM rep ()) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \(HistOp Shape
dest_shape SubExp
rf [VName]
dests [SubExp]
nes Lambda (Aliases rep)
op) -> do
    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)) dest_shape
    TC.require (Prim int64) rf

    -- Operator type must match the type of neutral elements.
    TC.checkLambda op $ 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 (nes_t == lambdaReturnType op) $
      TC.bad . TC.TypeError $
        "Operator has return type "
          <> prettyTuple (lambdaReturnType op)
          <> " but neutral element has type "
          <> prettyTuple nes_t

    -- Arrays must have proper type.
    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

  -- Types of input arrays must equal parameter types for bucket function.
  img' <- SubExp -> [VName] -> TypeM rep [Arg]
forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
TC.checkSOACArrayArgs SubExp
w [VName]
arrs
  TC.checkLambda bucket_fun img'

  -- Return type of bucket function must be an index for each
  -- operation followed by the values to write.
  nes_ts <- concat <$> mapM (mapM subExpType . histNeutral) ops
  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]
nes_ts
  unless (bucket_ret_t == lambdaReturnType bucket_fun) $
    TC.bad . TC.TypeError $
      "Bucket function has return type "
        <> prettyTuple (lambdaReturnType bucket_fun)
        <> " but should have type "
        <> prettyTuple bucket_ret_t
typeCheckSOAC (Screma SubExp
w [VName]
arrs (ScremaForm Lambda (Aliases rep)
map_lam [Scan (Aliases rep)]
scans [Reduce (Aliases rep)]
reds Lambda (Aliases rep)
post_lam)) = 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
w
  arrs' <- SubExp -> [VName] -> TypeM rep [Arg]
forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
TC.checkSOACArrayArgs SubExp
w [VName]
arrs
  TC.checkLambda map_lam arrs'
  scan_nes' <- concat <$> mapM typeCheckScan scans
  red_nes' <- concat <$> mapM typeCheckReduce reds
  let map_lam_ts = Lambda (Aliases rep) -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda (Aliases rep)
map_lam
  unless
    ( take (length scan_nes' + length red_nes') map_lam_ts
        == map TC.argType (scan_nes' ++ red_nes')
    )
    . TC.bad
    . TC.TypeError
    $ "Pre-lambda function return type "
      <> prettyTuple map_lam_ts
      <> " wrong for given scan and reduction functions."
  let (scan_ts, _, map_ts) =
        splitAt3 (length scan_nes') (length red_nes') map_lam_ts
      post_lam_args = (TypeBase Shape NoMode -> Arg) -> [TypeBase Shape NoMode] -> [Arg]
forall a b. (a -> b) -> [a] -> [b]
map (,Names
forall a. Monoid a => a
mempty) ([TypeBase Shape NoMode] -> [Arg])
-> [TypeBase Shape NoMode] -> [Arg]
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoMode]
scan_ts [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Semigroup a => a -> a -> a
<> [TypeBase Shape NoMode]
map_ts
  TC.checkLambda post_lam post_lam_args

  when (null scans && not (isIdentityLambda post_lam)) $
    TC.bad $
      TC.TypeError "Screma has post-lambda but no scan operations."

typeCheckScan :: (TC.Checkable rep) => Scan (Aliases rep) -> TC.TypeM rep [(Type, Names)]
typeCheckScan :: forall rep. Checkable rep => Scan (Aliases rep) -> TypeM rep [Arg]
typeCheckScan (Scan Lambda (Aliases rep)
scan_lam [SubExp]
scan_nes) = do
  scan_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]
scan_nes
  let scan_t = (Arg -> TypeBase Shape NoMode) -> [Arg] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> TypeBase Shape NoMode
TC.argType [Arg]
scan_nes'
  TC.checkLambda scan_lam $ map TC.noArgAliases $ scan_nes' ++ scan_nes'
  unless (scan_t == lambdaReturnType scan_lam) $
    TC.bad . TC.TypeError $
      "Scan function returns type "
        <> prettyTuple (lambdaReturnType scan_lam)
        <> " but neutral element has type "
        <> prettyTuple scan_t
  pure scan_nes'

typeCheckReduce :: (TC.Checkable rep) => Reduce (Aliases rep) -> TC.TypeM rep [(Type, Names)]
typeCheckReduce :: forall rep.
Checkable rep =>
Reduce (Aliases rep) -> TypeM rep [Arg]
typeCheckReduce (Reduce Commutativity
_ Lambda (Aliases rep)
red_lam [SubExp]
red_nes) = do
  red_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]
red_nes
  let red_t = (Arg -> TypeBase Shape NoMode) -> [Arg] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> TypeBase Shape NoMode
TC.argType [Arg]
red_nes'
  TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes'
  unless (red_t == lambdaReturnType red_lam) $
    TC.bad . TC.TypeError $
      "Reduce function returns type "
        <> prettyTuple (lambdaReturnType red_lam)
        <> " but neutral element has type "
        <> prettyTuple red_t
  pure red_nes'

instance RephraseOp SOAC where
  rephraseInOp :: forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> SOAC from -> m (SOAC to)
rephraseInOp Rephraser m from to
r (VJP Shape
shape [SubExp]
args [SubExp]
vec Lambda from
lam) =
    Shape -> [SubExp] -> [SubExp] -> Lambda to -> SOAC to
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
VJP Shape
shape [SubExp]
args [SubExp]
vec (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC 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
lam
  rephraseInOp Rephraser m from to
r (JVP Shape
shape [SubExp]
args [SubExp]
vec Lambda from
lam) =
    Shape -> [SubExp] -> [SubExp] -> Lambda to -> SOAC to
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
JVP Shape
shape [SubExp]
args [SubExp]
vec (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC 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
lam
  rephraseInOp Rephraser m from to
r (WithVJP [SubExp]
args Lambda from
lam Lambda from
lam_adj) =
    [SubExp] -> Lambda to -> Lambda to -> SOAC to
forall rep. [SubExp] -> Lambda rep -> Lambda rep -> SOAC rep
WithVJP [SubExp]
args (Lambda to -> Lambda to -> SOAC to)
-> m (Lambda to) -> m (Lambda to -> SOAC 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
lam m (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC 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 -> 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
lam_adj
  rephraseInOp Rephraser m from to
r (FlatMap SubExp
w [VName]
arrs ExtLambda from
lam) =
    SubExp -> [VName] -> ExtLambda to -> SOAC to
forall rep. SubExp -> [VName] -> ExtLambda rep -> SOAC rep
FlatMap SubExp
w [VName]
arrs (ExtLambda to -> SOAC to) -> m (ExtLambda to) -> m (SOAC to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser m from to -> ExtLambda from -> m (ExtLambda 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 ExtLambda from
lam
  rephraseInOp Rephraser m from to
r (Stream SubExp
w [VName]
arrs [SubExp]
acc Lambda from
lam) =
    SubExp -> [VName] -> [SubExp] -> Lambda to -> SOAC to
forall rep. SubExp -> [VName] -> [SubExp] -> Lambda rep -> SOAC rep
Stream SubExp
w [VName]
arrs [SubExp]
acc (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC 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
lam
  rephraseInOp Rephraser m from to
r (Hist SubExp
w [VName]
arrs [HistOp from]
ops Lambda from
lam) =
    SubExp -> [VName] -> [HistOp to] -> Lambda to -> SOAC to
forall rep.
SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
Hist SubExp
w [VName]
arrs ([HistOp to] -> Lambda to -> SOAC to)
-> m [HistOp to] -> m (Lambda to -> SOAC to)
forall (f :: * -> *) a b. Functor 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]
ops m (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC 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 -> 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
lam
    where
      onOp :: HistOp from -> m (HistOp to)
onOp (HistOp Shape
dest_shape SubExp
rf [VName]
dests [SubExp]
nes Lambda from
op) =
        Shape -> SubExp -> [VName] -> [SubExp] -> Lambda to -> HistOp to
forall rep.
Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
HistOp Shape
dest_shape SubExp
rf [VName]
dests [SubExp]
nes (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
  rephraseInOp Rephraser m from to
r (Screma SubExp
w [VName]
arrs (ScremaForm Lambda from
lam [Scan from]
scans [Reduce from]
red Lambda from
post_lam)) =
    SubExp -> [VName] -> ScremaForm to -> SOAC to
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs
      (ScremaForm to -> SOAC to) -> m (ScremaForm to) -> m (SOAC to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( Lambda to -> [Scan to] -> [Reduce to] -> Lambda to -> ScremaForm to
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm
              (Lambda to
 -> [Scan to] -> [Reduce to] -> Lambda to -> ScremaForm to)
-> m (Lambda to)
-> m ([Scan to] -> [Reduce to] -> Lambda to -> ScremaForm 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
lam
              m ([Scan to] -> [Reduce to] -> Lambda to -> ScremaForm to)
-> m [Scan to] -> m ([Reduce to] -> Lambda to -> ScremaForm to)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Scan from -> m (Scan to)) -> [Scan from] -> m [Scan 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 -> Scan from -> m (Scan to)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Scan from -> m (Scan to)
rephraseScan Rephraser m from to
r) [Scan from]
scans
              m ([Reduce to] -> Lambda to -> ScremaForm to)
-> m [Reduce to] -> m (Lambda to -> ScremaForm to)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Reduce from -> m (Reduce to)) -> [Reduce from] -> m [Reduce 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 -> Reduce from -> m (Reduce to)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Reduce from -> m (Reduce to)
rephraseRed Rephraser m from to
r) [Reduce from]
red
              m (Lambda to -> ScremaForm to)
-> m (Lambda to) -> m (ScremaForm 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 -> 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
post_lam
          )

rephraseRed :: (Monad m) => Rephraser m from to -> Reduce from -> m (Reduce to)
rephraseRed :: forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Reduce from -> m (Reduce to)
rephraseRed Rephraser m from to
r (Reduce Commutativity
comm Lambda from
op [SubExp]
nes) =
  Commutativity -> Lambda to -> [SubExp] -> Reduce to
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
comm (Lambda to -> [SubExp] -> Reduce to)
-> m (Lambda to) -> m ([SubExp] -> Reduce 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 m ([SubExp] -> Reduce to) -> m [SubExp] -> m (Reduce to)
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]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
nes

rephraseScan :: (Monad m) => Rephraser m from to -> Scan from -> m (Scan to)
rephraseScan :: forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Scan from -> m (Scan to)
rephraseScan Rephraser m from to
r (Scan Lambda from
op [SubExp]
nes) =
  Lambda to -> [SubExp] -> Scan to
forall rep. Lambda rep -> [SubExp] -> Scan rep
Scan (Lambda to -> [SubExp] -> Scan to)
-> m (Lambda to) -> m ([SubExp] -> Scan 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 m ([SubExp] -> Scan to) -> m [SubExp] -> m (Scan to)
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]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
nes

instance (OpMetrics (Op rep)) => OpMetrics (SOAC rep) where
  opMetrics :: SOAC rep -> MetricsM ()
opMetrics (VJP Shape
_ [SubExp]
_ [SubExp]
_ Lambda rep
lam) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"VJP" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
lam
  opMetrics (JVP Shape
_ [SubExp]
_ [SubExp]
_ Lambda rep
lam) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"JVP" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
lam
  opMetrics (WithVJP [SubExp]
_ Lambda rep
lam Lambda rep
lam_adj) = do
    Text -> MetricsM () -> MetricsM ()
inside Text
"WithVJP" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
lam
    Text -> MetricsM () -> MetricsM ()
inside Text
"WithVJP" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
lam_adj
  opMetrics (FlatMap SubExp
_ [VName]
_ ExtLambda rep
lam) = do
    Text -> MetricsM () -> MetricsM ()
inside Text
"FlatMap" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ ExtLambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics ExtLambda rep
lam
  opMetrics (Stream SubExp
_ [VName]
_ [SubExp]
_ Lambda rep
lam) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"Stream" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
lam
  opMetrics (Hist SubExp
_ [VName]
_ [HistOp rep]
ops Lambda rep
bucket_fun) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"Hist" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ (HistOp rep -> MetricsM ()) -> [HistOp rep] -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics (Lambda rep -> MetricsM ())
-> (HistOp rep -> Lambda rep) -> 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 -> Lambda rep
forall rep. HistOp rep -> Lambda rep
histOp) [HistOp rep]
ops MetricsM () -> MetricsM () -> MetricsM ()
forall a b. MetricsM a -> MetricsM b -> MetricsM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
bucket_fun
  opMetrics (Screma SubExp
_ [VName]
_ (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam)) =
    Text -> MetricsM () -> MetricsM ()
inside Text
"Screma" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ do
      Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
map_lam
      (Scan rep -> MetricsM ()) -> [Scan rep] -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics (Lambda rep -> MetricsM ())
-> (Scan rep -> Lambda rep) -> Scan 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
. Scan rep -> Lambda rep
forall rep. Scan rep -> Lambda rep
scanLambda) [Scan rep]
scans
      (Reduce rep -> MetricsM ()) -> [Reduce rep] -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics (Lambda rep -> MetricsM ())
-> (Reduce rep -> Lambda rep) -> Reduce 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
. Reduce rep -> Lambda rep
forall rep. Reduce rep -> Lambda rep
redLambda) [Reduce rep]
reds
      Lambda rep -> MetricsM ()
forall rep t. OpMetrics (Op rep) => GLambda rep t -> MetricsM ()
lambdaMetrics Lambda rep
post_lam

instance (PrettyRep rep) => PP.Pretty (SOAC rep) where
  pretty :: forall ann. SOAC rep -> Doc ann
pretty (VJP Shape
shape [SubExp]
args [SubExp]
vec Lambda rep
lam) =
    Doc ann
"vjp"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens
        ( Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.align (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
            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
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
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]
args)
              Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
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]
vec)
              Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
lam
        )
  pretty (JVP Shape
shape [SubExp]
args [SubExp]
vec Lambda rep
lam) =
    Doc ann
"jvp"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens
        ( Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.align (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
            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
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
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]
args)
              Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
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]
vec)
              Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
lam
        )
  pretty (WithVJP [SubExp]
args Lambda rep
lam Lambda rep
lam_adj) =
    Doc ann
"with_vjp"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens
        ( Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.align (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
            Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.braces ([Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
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]
args)
              Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
lam
              Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
lam_adj
        )
  pretty (FlatMap SubExp
w [VName]
arrs ExtLambda rep
lam) =
    Doc ann
"flatmap"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
        ( SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
w
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma 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' ((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]
arrs)
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> ExtLambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. ExtLambda rep -> Doc ann
pretty ExtLambda rep
lam
        )
  pretty (Stream SubExp
size [VName]
arrs [SubExp]
acc Lambda rep
lam) =
    SubExp -> [VName] -> [SubExp] -> Lambda rep -> Doc ann
forall rep inp ann.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> [SubExp] -> Lambda rep -> Doc ann
ppStream SubExp
size [VName]
arrs [SubExp]
acc Lambda rep
lam
  pretty (Hist SubExp
w [VName]
arrs [HistOp rep]
ops Lambda rep
bucket_fun) =
    SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> Doc ann
forall rep inp ann.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> [HistOp rep] -> Lambda rep -> Doc ann
ppHist SubExp
w [VName]
arrs [HistOp rep]
ops Lambda rep
bucket_fun
  pretty (Screma SubExp
w [VName]
arrs ScremaForm rep
screma)
    | Just Lambda rep
map_lam <- ScremaForm rep -> Maybe (Lambda rep)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm rep
screma =
        Doc ann
"map"
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
            ( SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
w
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma 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' ((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]
arrs)
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
map_lam
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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 (ScremaForm rep -> Lambda rep
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm rep
screma)
            )
    | Just ([Reduce rep]
reds, Lambda rep
map_lam) <- ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm rep
screma =
        Doc ann
"redomap"
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
            ( SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
w
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma 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' ((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]
arrs)
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
map_lam
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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. 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
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
$ (Reduce rep -> Doc ann) -> [Reduce rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Reduce rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Reduce rep -> Doc ann
pretty [Reduce rep]
reds)
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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 (ScremaForm rep -> Lambda rep
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm rep
screma)
            )
    | Just ([Scan rep]
scans, Lambda rep
map_lam) <- ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
isScanomapSOAC ScremaForm rep
screma =
        Doc ann
"scanomap"
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
            ( SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
w
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma 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' ((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]
arrs)
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
map_lam
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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. 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
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
$ (Scan rep -> Doc ann) -> [Scan rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Scan rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Scan rep -> Doc ann
pretty [Scan rep]
scans)
                Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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 (ScremaForm rep -> Lambda rep
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm rep
screma)
            )
  pretty (Screma SubExp
w [VName]
arrs ScremaForm rep
form) = SubExp -> [VName] -> ScremaForm rep -> Doc ann
forall rep inp ann.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> ScremaForm rep -> Doc ann
ppScrema SubExp
w [VName]
arrs ScremaForm rep
form

-- | Prettyprint the given Screma.
ppScrema ::
  (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> ScremaForm rep -> Doc ann
ppScrema :: forall rep inp ann.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> ScremaForm rep -> Doc ann
ppScrema SubExp
w [inp]
arrs (ScremaForm Lambda rep
map_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam) =
  Doc ann
"screma"
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
      ( SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
w
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma 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' ((inp -> Doc ann) -> [inp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map inp -> Doc ann
forall ann. inp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [inp]
arrs)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
map_lam
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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. 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
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
$ (Scan rep -> Doc ann) -> [Scan rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Scan rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Scan rep -> Doc ann
pretty [Scan rep]
scans)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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. 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
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
$ (Reduce rep -> Doc ann) -> [Reduce rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Reduce rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Reduce rep -> Doc ann
pretty [Reduce rep]
reds)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
post_lam
      )

-- | Prettyprint the given Stream.
ppStream ::
  (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> [SubExp] -> Lambda rep -> Doc ann
ppStream :: forall rep inp ann.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> [SubExp] -> Lambda rep -> Doc ann
ppStream SubExp
size [inp]
arrs [SubExp]
acc Lambda rep
lam =
  Doc ann
"streamSeq"
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
      ( SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
size
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
            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' ((inp -> Doc ann) -> [inp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map inp -> Doc ann
forall ann. inp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [inp]
arrs)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
            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' ((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]
acc)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
lam
      )

instance (PrettyRep rep) => Pretty (Scan rep) where
  pretty :: forall ann. Scan rep -> Doc ann
pretty (Scan Lambda rep
scan_lam [SubExp]
scan_nes) =
    Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
scan_lam Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
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]
scan_nes)

ppComm :: Commutativity -> Doc ann
ppComm :: forall ann. Commutativity -> Doc ann
ppComm Commutativity
Noncommutative = Doc ann
forall a. Monoid a => a
mempty
ppComm Commutativity
Commutative = Doc ann
"commutative "

instance (PrettyRep rep) => Pretty (Reduce rep) where
  pretty :: forall ann. Reduce rep -> Doc ann
pretty (Reduce Commutativity
comm Lambda rep
red_lam [SubExp]
red_nes) =
    Commutativity -> Doc ann
forall ann. Commutativity -> Doc ann
ppComm Commutativity
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
red_lam
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
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]
red_nes)

-- | Prettyprint the given histogram operation.
ppHist ::
  (PrettyRep rep, Pretty inp) =>
  SubExp ->
  [inp] ->
  [HistOp rep] ->
  Lambda rep ->
  Doc ann
ppHist :: forall rep inp ann.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> [HistOp rep] -> Lambda rep -> Doc ann
ppHist SubExp
w [inp]
arrs [HistOp rep]
ops Lambda rep
bucket_fun =
  Doc ann
"hist"
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
      ( SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
w
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
            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' ((inp -> Doc ann) -> [inp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map inp -> Doc ann
forall ann. inp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [inp]
arrs)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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. 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
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)
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
bucket_fun
      )
  where
    ppOp :: HistOp rep -> Doc ann
ppOp (HistOp Shape
dest_w SubExp
rf [VName]
dests [SubExp]
nes Lambda rep
op) =
      Shape -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Shape -> Doc ann
pretty Shape
dest_w
        Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
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
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
comma
          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' ((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
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 (PrettyRep rep) => PP.Pretty (ScremaForm rep) where
  pretty :: forall ann. ScremaForm rep -> Doc ann
pretty (ScremaForm Lambda rep
pre_lam [Scan rep]
scans [Reduce rep]
reds Lambda rep
post_lam) =
    Doc ann
"screma"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
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
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
        ( Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
pre_lam
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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. 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
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
$ (Scan rep -> Doc ann) -> [Scan rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Scan rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Scan rep -> Doc ann
pretty [Scan rep]
scans)
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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. 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
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
$ (Reduce rep -> Doc ann) -> [Reduce rep] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Reduce rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Reduce rep -> Doc ann
pretty [Reduce rep]
reds)
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
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
post_lam
        )