module Futhark.Pass.Flatten.Distribute
  ( distributeMap,
    distributeBody,
    MapArray (..),
    mapArrayRowType,
    DistResults (..),
    DistRep,
    ResMap,
    Distributed (..),
    DistStm (..),
    DistStms,
    DistBody (..),
    DistInput (..),
    DistInputs,
    DistType (..),
    distInputType,
    DistResult (..),
    ResTag (..),
    DistIrregularity (..),
    FunHasParallelism,
    isRegularDistResult,
    isParallelStm,
    stmHasMeaningfulParallelism,

    -- * Segments
    Segments,
    segmentsShape,
    segmentsRank,
    segmentCount,
  )
where

import Data.Bifunctor
import Data.Foldable
import Data.List qualified as L
import Data.Map qualified as M
import Data.Maybe
import Data.Sequence qualified as Seq
import Data.Set qualified as S
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.SOACS
import Futhark.Util (nubOrd)
import Futhark.Util.Pretty

-- | Widths of the enclosing map-nest, outermost first. For top-level parallel
-- constructs, this is empty, which should be treated as an implicit
-- single-element segment (see 'segmentCount'). Generally, the empty-segments
-- case must be treated specially in some places, which is unfortunate, but it
-- helps unify code between the nested and top level cases.
type Segments = [SubExp]

type FunHasParallelism = Name -> Bool

-- | How to treat irregularity when classifying the statements of a distributed
-- body. This is mainly used to sequentialise nonuniform nested parallelism
-- instead of actually exploiting the parallelism, as the overhead of doing so
-- can sometimes be ruinous.
data DistIrregularity
  = -- | Distribute statements involving irregularity, relying on the machinery
    -- for flattening irregular arrays to handle them.
    DistributeIrregular
  | -- | Sequentialise BasicOps that involve only internal nonuniformity instead
    -- of distributing them. Used when generating intrablock code, where the
    -- machinery for flattening nonuniform nested parallelism would produce
    -- SegOps whose sizes are bound inside the kernel body, which makes the
    -- enclosing intrablock kernel infeasible ('noNonuniformPar' would reject
    -- it). Irregularity that escapes the enclosing map must still be
    -- distributed; if it occurs, the intrablock version is correctly rejected.
    SequentialiseIrregularBasicOps
  | -- | Sequentialise /any/ statement whose nonuniformity stays internal.
    SequentialiseIrregularAll
  deriving (DistIrregularity -> DistIrregularity -> Bool
(DistIrregularity -> DistIrregularity -> Bool)
-> (DistIrregularity -> DistIrregularity -> Bool)
-> Eq DistIrregularity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DistIrregularity -> DistIrregularity -> Bool
== :: DistIrregularity -> DistIrregularity -> Bool
$c/= :: DistIrregularity -> DistIrregularity -> Bool
/= :: DistIrregularity -> DistIrregularity -> Bool
Eq, Int -> DistIrregularity -> ShowS
[DistIrregularity] -> ShowS
DistIrregularity -> String
(Int -> DistIrregularity -> ShowS)
-> (DistIrregularity -> String)
-> ([DistIrregularity] -> ShowS)
-> Show DistIrregularity
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DistIrregularity -> ShowS
showsPrec :: Int -> DistIrregularity -> ShowS
$cshow :: DistIrregularity -> String
show :: DistIrregularity -> String
$cshowList :: [DistIrregularity] -> ShowS
showList :: [DistIrregularity] -> ShowS
Show)

segmentsShape :: Segments -> Shape
segmentsShape :: Segments -> Shape
segmentsShape = Segments -> Shape
forall d. [d] -> ShapeBase d
Shape

segmentsRank :: Segments -> Int
segmentsRank :: Segments -> Int
segmentsRank = Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> (Segments -> Shape) -> Segments -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Segments -> Shape
segmentsShape

segmentCount :: Segments -> TPrimExp Int64 VName
segmentCount :: Segments -> TPrimExp Int64 VName
segmentCount = [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
-> (Segments -> [TPrimExp Int64 VName])
-> Segments
-> TPrimExp Int64 VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> TPrimExp Int64 VName)
-> Segments -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 (Segments -> [TPrimExp Int64 VName])
-> (Segments -> Segments) -> Segments -> [TPrimExp Int64 VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> Segments
forall d. ShapeBase d -> [d]
shapeDims (Shape -> Segments) -> (Segments -> Shape) -> Segments -> Segments
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Segments -> Shape
segmentsShape

newtype ResTag = ResTag Int
  deriving (ResTag -> ResTag -> Bool
(ResTag -> ResTag -> Bool)
-> (ResTag -> ResTag -> Bool) -> Eq ResTag
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ResTag -> ResTag -> Bool
== :: ResTag -> ResTag -> Bool
$c/= :: ResTag -> ResTag -> Bool
/= :: ResTag -> ResTag -> Bool
Eq, Eq ResTag
Eq ResTag =>
(ResTag -> ResTag -> Ordering)
-> (ResTag -> ResTag -> Bool)
-> (ResTag -> ResTag -> Bool)
-> (ResTag -> ResTag -> Bool)
-> (ResTag -> ResTag -> Bool)
-> (ResTag -> ResTag -> ResTag)
-> (ResTag -> ResTag -> ResTag)
-> Ord ResTag
ResTag -> ResTag -> Bool
ResTag -> ResTag -> Ordering
ResTag -> ResTag -> ResTag
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: ResTag -> ResTag -> Ordering
compare :: ResTag -> ResTag -> Ordering
$c< :: ResTag -> ResTag -> Bool
< :: ResTag -> ResTag -> Bool
$c<= :: ResTag -> ResTag -> Bool
<= :: ResTag -> ResTag -> Bool
$c> :: ResTag -> ResTag -> Bool
> :: ResTag -> ResTag -> Bool
$c>= :: ResTag -> ResTag -> Bool
>= :: ResTag -> ResTag -> Bool
$cmax :: ResTag -> ResTag -> ResTag
max :: ResTag -> ResTag -> ResTag
$cmin :: ResTag -> ResTag -> ResTag
min :: ResTag -> ResTag -> ResTag
Ord, Int -> ResTag -> ShowS
[ResTag] -> ShowS
ResTag -> String
(Int -> ResTag -> ShowS)
-> (ResTag -> String) -> ([ResTag] -> ShowS) -> Show ResTag
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ResTag -> ShowS
showsPrec :: Int -> ResTag -> ShowS
$cshow :: ResTag -> String
show :: ResTag -> String
$cshowList :: [ResTag] -> ShowS
showList :: [ResTag] -> ShowS
Show)

-- | Something that is mapped.
data DistInput
  = -- | A value bound outside the original map-nest.  By necessity
    -- regular.  The type is the parameter type.
    DistInputFree VName Type
  | -- | A value constructed inside the original map-nest.  May be
    -- irregular.
    DistInput ResTag Type
  deriving (DistInput -> DistInput -> Bool
(DistInput -> DistInput -> Bool)
-> (DistInput -> DistInput -> Bool) -> Eq DistInput
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DistInput -> DistInput -> Bool
== :: DistInput -> DistInput -> Bool
$c/= :: DistInput -> DistInput -> Bool
/= :: DistInput -> DistInput -> Bool
Eq, Eq DistInput
Eq DistInput =>
(DistInput -> DistInput -> Ordering)
-> (DistInput -> DistInput -> Bool)
-> (DistInput -> DistInput -> Bool)
-> (DistInput -> DistInput -> Bool)
-> (DistInput -> DistInput -> Bool)
-> (DistInput -> DistInput -> DistInput)
-> (DistInput -> DistInput -> DistInput)
-> Ord DistInput
DistInput -> DistInput -> Bool
DistInput -> DistInput -> Ordering
DistInput -> DistInput -> DistInput
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DistInput -> DistInput -> Ordering
compare :: DistInput -> DistInput -> Ordering
$c< :: DistInput -> DistInput -> Bool
< :: DistInput -> DistInput -> Bool
$c<= :: DistInput -> DistInput -> Bool
<= :: DistInput -> DistInput -> Bool
$c> :: DistInput -> DistInput -> Bool
> :: DistInput -> DistInput -> Bool
$c>= :: DistInput -> DistInput -> Bool
>= :: DistInput -> DistInput -> Bool
$cmax :: DistInput -> DistInput -> DistInput
max :: DistInput -> DistInput -> DistInput
$cmin :: DistInput -> DistInput -> DistInput
min :: DistInput -> DistInput -> DistInput
Ord, Int -> DistInput -> ShowS
[DistInput] -> ShowS
DistInput -> String
(Int -> DistInput -> ShowS)
-> (DistInput -> String)
-> ([DistInput] -> ShowS)
-> Show DistInput
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DistInput -> ShowS
showsPrec :: Int -> DistInput -> ShowS
$cshow :: DistInput -> String
show :: DistInput -> String
$cshowList :: [DistInput] -> ShowS
showList :: [DistInput] -> ShowS
Show)

type DistInputs = [(VName, DistInput)]

nubInputs :: DistInputs -> DistInputs
nubInputs :: DistInputs -> DistInputs
nubInputs = ((VName, DistInput) -> (VName, DistInput) -> Bool)
-> DistInputs -> DistInputs
forall a. (a -> a -> Bool) -> [a] -> [a]
L.nubBy (\(VName, DistInput)
a (VName, DistInput)
b -> (VName, DistInput) -> VName
forall a b. (a, b) -> a
fst (VName, DistInput)
a VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== (VName, DistInput) -> VName
forall a b. (a, b) -> a
fst (VName, DistInput)
b)

-- | The type of a 'DistInput'.  This corresponds to the parameter
-- type of the original map-nest.
distInputType :: DistInput -> Type
distInputType :: DistInput -> Type
distInputType (DistInputFree VName
_ Type
t) = Type
t
distInputType (DistInput ResTag
_ Type
t) = Type
t

data DistType
  = DistType
      -- | Outer regular size.
      Segments
      -- | Irregular dimensions on top (but after the leading regular
      -- size).
      Rank
      -- | The regular "element type" - in the worst case, at least a
      -- scalar.
      Type
  deriving (DistType -> DistType -> Bool
(DistType -> DistType -> Bool)
-> (DistType -> DistType -> Bool) -> Eq DistType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DistType -> DistType -> Bool
== :: DistType -> DistType -> Bool
$c/= :: DistType -> DistType -> Bool
/= :: DistType -> DistType -> Bool
Eq, Eq DistType
Eq DistType =>
(DistType -> DistType -> Ordering)
-> (DistType -> DistType -> Bool)
-> (DistType -> DistType -> Bool)
-> (DistType -> DistType -> Bool)
-> (DistType -> DistType -> Bool)
-> (DistType -> DistType -> DistType)
-> (DistType -> DistType -> DistType)
-> Ord DistType
DistType -> DistType -> Bool
DistType -> DistType -> Ordering
DistType -> DistType -> DistType
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DistType -> DistType -> Ordering
compare :: DistType -> DistType -> Ordering
$c< :: DistType -> DistType -> Bool
< :: DistType -> DistType -> Bool
$c<= :: DistType -> DistType -> Bool
<= :: DistType -> DistType -> Bool
$c> :: DistType -> DistType -> Bool
> :: DistType -> DistType -> Bool
$c>= :: DistType -> DistType -> Bool
>= :: DistType -> DistType -> Bool
$cmax :: DistType -> DistType -> DistType
max :: DistType -> DistType -> DistType
$cmin :: DistType -> DistType -> DistType
min :: DistType -> DistType -> DistType
Ord, Int -> DistType -> ShowS
[DistType] -> ShowS
DistType -> String
(Int -> DistType -> ShowS)
-> (DistType -> String) -> ([DistType] -> ShowS) -> Show DistType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DistType -> ShowS
showsPrec :: Int -> DistType -> ShowS
$cshow :: DistType -> String
show :: DistType -> String
$cshowList :: [DistType] -> ShowS
showList :: [DistType] -> ShowS
Show)

data DistResult = DistResult {DistResult -> ResTag
distResTag :: ResTag, DistResult -> DistType
distResType :: DistType, DistResult -> VName
distResName :: VName}
  deriving (DistResult -> DistResult -> Bool
(DistResult -> DistResult -> Bool)
-> (DistResult -> DistResult -> Bool) -> Eq DistResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DistResult -> DistResult -> Bool
== :: DistResult -> DistResult -> Bool
$c/= :: DistResult -> DistResult -> Bool
/= :: DistResult -> DistResult -> Bool
Eq, Eq DistResult
Eq DistResult =>
(DistResult -> DistResult -> Ordering)
-> (DistResult -> DistResult -> Bool)
-> (DistResult -> DistResult -> Bool)
-> (DistResult -> DistResult -> Bool)
-> (DistResult -> DistResult -> Bool)
-> (DistResult -> DistResult -> DistResult)
-> (DistResult -> DistResult -> DistResult)
-> Ord DistResult
DistResult -> DistResult -> Bool
DistResult -> DistResult -> Ordering
DistResult -> DistResult -> DistResult
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DistResult -> DistResult -> Ordering
compare :: DistResult -> DistResult -> Ordering
$c< :: DistResult -> DistResult -> Bool
< :: DistResult -> DistResult -> Bool
$c<= :: DistResult -> DistResult -> Bool
<= :: DistResult -> DistResult -> Bool
$c> :: DistResult -> DistResult -> Bool
> :: DistResult -> DistResult -> Bool
$c>= :: DistResult -> DistResult -> Bool
>= :: DistResult -> DistResult -> Bool
$cmax :: DistResult -> DistResult -> DistResult
max :: DistResult -> DistResult -> DistResult
$cmin :: DistResult -> DistResult -> DistResult
min :: DistResult -> DistResult -> DistResult
Ord, Int -> DistResult -> ShowS
[DistResult] -> ShowS
DistResult -> String
(Int -> DistResult -> ShowS)
-> (DistResult -> String)
-> ([DistResult] -> ShowS)
-> Show DistResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DistResult -> ShowS
showsPrec :: Int -> DistResult -> ShowS
$cshow :: DistResult -> String
show :: DistResult -> String
$cshowList :: [DistResult] -> ShowS
showList :: [DistResult] -> ShowS
Show)

-- | The body of a distributed statement.
data DistBody
  = -- | A single statement that may involve parallel operations or produce an
    -- irregular array.
    ParallelStm (Stm SOACS)
  | -- | Single or Multiple scalar operations grouped into a single traversal
    ScalarStm (Stms SOACS)
  deriving (DistBody -> DistBody -> Bool
(DistBody -> DistBody -> Bool)
-> (DistBody -> DistBody -> Bool) -> Eq DistBody
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DistBody -> DistBody -> Bool
== :: DistBody -> DistBody -> Bool
$c/= :: DistBody -> DistBody -> Bool
/= :: DistBody -> DistBody -> Bool
Eq, Eq DistBody
Eq DistBody =>
(DistBody -> DistBody -> Ordering)
-> (DistBody -> DistBody -> Bool)
-> (DistBody -> DistBody -> Bool)
-> (DistBody -> DistBody -> Bool)
-> (DistBody -> DistBody -> Bool)
-> (DistBody -> DistBody -> DistBody)
-> (DistBody -> DistBody -> DistBody)
-> Ord DistBody
DistBody -> DistBody -> Bool
DistBody -> DistBody -> Ordering
DistBody -> DistBody -> DistBody
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DistBody -> DistBody -> Ordering
compare :: DistBody -> DistBody -> Ordering
$c< :: DistBody -> DistBody -> Bool
< :: DistBody -> DistBody -> Bool
$c<= :: DistBody -> DistBody -> Bool
<= :: DistBody -> DistBody -> Bool
$c> :: DistBody -> DistBody -> Bool
> :: DistBody -> DistBody -> Bool
$c>= :: DistBody -> DistBody -> Bool
>= :: DistBody -> DistBody -> Bool
$cmax :: DistBody -> DistBody -> DistBody
max :: DistBody -> DistBody -> DistBody
$cmin :: DistBody -> DistBody -> DistBody
min :: DistBody -> DistBody -> DistBody
Ord, Int -> DistBody -> ShowS
[DistBody] -> ShowS
DistBody -> String
(Int -> DistBody -> ShowS)
-> (DistBody -> String) -> ([DistBody] -> ShowS) -> Show DistBody
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DistBody -> ShowS
showsPrec :: Int -> DistBody -> ShowS
$cshow :: DistBody -> String
show :: DistBody -> String
$cshowList :: [DistBody] -> ShowS
showList :: [DistBody] -> ShowS
Show)

distBodyStms :: DistBody -> Stms SOACS
distBodyStms :: DistBody -> Stms SOACS
distBodyStms (ParallelStm Stm SOACS
stm) = Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm
distBodyStms (ScalarStm Stms SOACS
stms) = Stms SOACS
stms

data DistStm = DistStm
  { DistStm -> DistInputs
distStmInputs :: DistInputs,
    DistStm -> [DistResult]
distStmResult :: [DistResult],
    DistStm -> DistBody
distStmBody :: DistBody
  }
  deriving (DistStm -> DistStm -> Bool
(DistStm -> DistStm -> Bool)
-> (DistStm -> DistStm -> Bool) -> Eq DistStm
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DistStm -> DistStm -> Bool
== :: DistStm -> DistStm -> Bool
$c/= :: DistStm -> DistStm -> Bool
/= :: DistStm -> DistStm -> Bool
Eq, Eq DistStm
Eq DistStm =>
(DistStm -> DistStm -> Ordering)
-> (DistStm -> DistStm -> Bool)
-> (DistStm -> DistStm -> Bool)
-> (DistStm -> DistStm -> Bool)
-> (DistStm -> DistStm -> Bool)
-> (DistStm -> DistStm -> DistStm)
-> (DistStm -> DistStm -> DistStm)
-> Ord DistStm
DistStm -> DistStm -> Bool
DistStm -> DistStm -> Ordering
DistStm -> DistStm -> DistStm
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DistStm -> DistStm -> Ordering
compare :: DistStm -> DistStm -> Ordering
$c< :: DistStm -> DistStm -> Bool
< :: DistStm -> DistStm -> Bool
$c<= :: DistStm -> DistStm -> Bool
<= :: DistStm -> DistStm -> Bool
$c> :: DistStm -> DistStm -> Bool
> :: DistStm -> DistStm -> Bool
$c>= :: DistStm -> DistStm -> Bool
>= :: DistStm -> DistStm -> Bool
$cmax :: DistStm -> DistStm -> DistStm
max :: DistStm -> DistStm -> DistStm
$cmin :: DistStm -> DistStm -> DistStm
min :: DistStm -> DistStm -> DistStm
Ord, Int -> DistStm -> ShowS
[DistStm] -> ShowS
DistStm -> String
(Int -> DistStm -> ShowS)
-> (DistStm -> String) -> ([DistStm] -> ShowS) -> Show DistStm
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DistStm -> ShowS
showsPrec :: Int -> DistStm -> ShowS
$cshow :: DistStm -> String
show :: DistStm -> String
$cshowList :: [DistStm] -> ShowS
showList :: [DistStm] -> ShowS
Show)

distStmStms :: DistStm -> Stms SOACS
distStmStms :: DistStm -> Stms SOACS
distStmStms = DistBody -> Stms SOACS
distBodyStms (DistBody -> Stms SOACS)
-> (DistStm -> DistBody) -> DistStm -> Stms SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistStm -> DistBody
distStmBody

-- | An efficient sequence of 'DistStm's.
type DistStms = Seq.Seq DistStm

-- | First element of tuple are certificates for this result.
--
-- Second is the name to which is should be bound.
--
-- Third is the element type (i.e. excluding shape of segments).
type ResMap = M.Map ResTag [([DistInput], VName, Type)]

-- | The results of a map-distribution that were free or identity
-- mapped in the original map function.  These correspond to plain
-- replicated arrays.
type DistRep = (VName, Either SubExp DistInput)

data DistResults = DistResults ResMap [DistRep]
  deriving (DistResults -> DistResults -> Bool
(DistResults -> DistResults -> Bool)
-> (DistResults -> DistResults -> Bool) -> Eq DistResults
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DistResults -> DistResults -> Bool
== :: DistResults -> DistResults -> Bool
$c/= :: DistResults -> DistResults -> Bool
/= :: DistResults -> DistResults -> Bool
Eq, Eq DistResults
Eq DistResults =>
(DistResults -> DistResults -> Ordering)
-> (DistResults -> DistResults -> Bool)
-> (DistResults -> DistResults -> Bool)
-> (DistResults -> DistResults -> Bool)
-> (DistResults -> DistResults -> Bool)
-> (DistResults -> DistResults -> DistResults)
-> (DistResults -> DistResults -> DistResults)
-> Ord DistResults
DistResults -> DistResults -> Bool
DistResults -> DistResults -> Ordering
DistResults -> DistResults -> DistResults
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DistResults -> DistResults -> Ordering
compare :: DistResults -> DistResults -> Ordering
$c< :: DistResults -> DistResults -> Bool
< :: DistResults -> DistResults -> Bool
$c<= :: DistResults -> DistResults -> Bool
<= :: DistResults -> DistResults -> Bool
$c> :: DistResults -> DistResults -> Bool
> :: DistResults -> DistResults -> Bool
$c>= :: DistResults -> DistResults -> Bool
>= :: DistResults -> DistResults -> Bool
$cmax :: DistResults -> DistResults -> DistResults
max :: DistResults -> DistResults -> DistResults
$cmin :: DistResults -> DistResults -> DistResults
min :: DistResults -> DistResults -> DistResults
Ord, Int -> DistResults -> ShowS
[DistResults] -> ShowS
DistResults -> String
(Int -> DistResults -> ShowS)
-> (DistResults -> String)
-> ([DistResults] -> ShowS)
-> Show DistResults
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DistResults -> ShowS
showsPrec :: Int -> DistResults -> ShowS
$cshow :: DistResults -> String
show :: DistResults -> String
$cshowList :: [DistResults] -> ShowS
showList :: [DistResults] -> ShowS
Show)

data Distributed = Distributed DistStms DistResults
  deriving (Distributed -> Distributed -> Bool
(Distributed -> Distributed -> Bool)
-> (Distributed -> Distributed -> Bool) -> Eq Distributed
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Distributed -> Distributed -> Bool
== :: Distributed -> Distributed -> Bool
$c/= :: Distributed -> Distributed -> Bool
/= :: Distributed -> Distributed -> Bool
Eq, Eq Distributed
Eq Distributed =>
(Distributed -> Distributed -> Ordering)
-> (Distributed -> Distributed -> Bool)
-> (Distributed -> Distributed -> Bool)
-> (Distributed -> Distributed -> Bool)
-> (Distributed -> Distributed -> Bool)
-> (Distributed -> Distributed -> Distributed)
-> (Distributed -> Distributed -> Distributed)
-> Ord Distributed
Distributed -> Distributed -> Bool
Distributed -> Distributed -> Ordering
Distributed -> Distributed -> Distributed
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Distributed -> Distributed -> Ordering
compare :: Distributed -> Distributed -> Ordering
$c< :: Distributed -> Distributed -> Bool
< :: Distributed -> Distributed -> Bool
$c<= :: Distributed -> Distributed -> Bool
<= :: Distributed -> Distributed -> Bool
$c> :: Distributed -> Distributed -> Bool
> :: Distributed -> Distributed -> Bool
$c>= :: Distributed -> Distributed -> Bool
>= :: Distributed -> Distributed -> Bool
$cmax :: Distributed -> Distributed -> Distributed
max :: Distributed -> Distributed -> Distributed
$cmin :: Distributed -> Distributed -> Distributed
min :: Distributed -> Distributed -> Distributed
Ord, Int -> Distributed -> ShowS
[Distributed] -> ShowS
Distributed -> String
(Int -> Distributed -> ShowS)
-> (Distributed -> String)
-> ([Distributed] -> ShowS)
-> Show Distributed
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Distributed -> ShowS
showsPrec :: Int -> Distributed -> ShowS
$cshow :: Distributed -> String
show :: Distributed -> String
$cshowList :: [Distributed] -> ShowS
showList :: [Distributed] -> ShowS
Show)

instance Pretty ResTag where
  pretty :: forall ann. ResTag -> Doc ann
pretty (ResTag Int
x) = Doc ann
"r" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Int
x

instance Pretty DistInput where
  pretty :: forall ann. DistInput -> Doc ann
pretty (DistInputFree VName
v Type
_) = VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
v
  pretty (DistInput ResTag
rt Type
_) = ResTag -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. ResTag -> Doc ann
pretty ResTag
rt

instance Pretty DistType where
  pretty :: forall ann. DistType -> Doc ann
pretty (DistType Segments
w Rank
r Type
t) =
    Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
brackets (Segments -> Doc ann
forall ann. Segments -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Segments
w) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Rank -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Rank -> Doc ann
pretty Rank
r Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Type -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Type -> Doc ann
pretty Type
t

instance Pretty DistResult where
  pretty :: forall ann. DistResult -> Doc ann
pretty (DistResult ResTag
rt DistType
t VName
_) =
    ResTag -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. ResTag -> Doc ann
pretty ResTag
rt Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> DistType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. DistType -> Doc ann
pretty DistType
t

instance Pretty DistStm where
  pretty :: forall ann. DistStm -> Doc ann
pretty (DistStm DistInputs
inputs [DistResult]
res DistBody
stms) =
    Doc ann
"let" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
ppTuple' ((DistResult -> Doc ann) -> [DistResult] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map DistResult -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. DistResult -> Doc ann
pretty [DistResult]
res) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"=" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 Doc ann
forall ann. Doc ann
stm'
    where
      stm' :: Doc ann
stm' =
        Doc ann
"map"
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
nestedBlock
            ( [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
stack ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$
                ((VName, DistInput) -> Doc ann) -> DistInputs -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (VName, DistInput) -> Doc ann
forall {a} {ann}. Pretty a => (a, DistInput) -> Doc ann
onInput DistInputs
inputs
                  [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. [a] -> [a] -> [a]
++ (Stm SOACS -> Doc ann) -> [Stm SOACS] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Stm SOACS -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Stm SOACS -> Doc ann
pretty (Stms SOACS -> [Stm SOACS]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (DistBody -> Stms SOACS
distBodyStms DistBody
stms))
                  [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. [a] -> [a] -> [a]
++ [ Doc ann
"return" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
ppTuple' ((DistResult -> Doc ann) -> [DistResult] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map DistResult -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. DistResult -> Doc ann
pretty [DistResult]
res)
                     ]
            )
      onInput :: (a, DistInput) -> Doc ann
onInput (a
v, DistInput
inp) =
        Doc ann
"for"
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
v Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Type -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Type -> Doc ann
pretty (DistInput -> Type
distInputType DistInput
inp))
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"<-"
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> DistInput -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. DistInput -> Doc ann
pretty DistInput
inp

instance Pretty Distributed where
  pretty :: forall ann. Distributed -> Doc ann
pretty (Distributed Seq DistStm
stms (DistResults ResMap
resmap [DistRep]
reps)) =
    Doc ann
forall ann. Doc ann
stms' Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ann
forall ann. Doc ann
res'
    where
      res' :: Doc a
res' = [Doc a] -> Doc a
forall a. [Doc a] -> Doc a
stack ([Doc a] -> Doc a) -> [Doc a] -> Doc a
forall a b. (a -> b) -> a -> b
$ ((ResTag, [([DistInput], VName, Type)]) -> Doc a)
-> [(ResTag, [([DistInput], VName, Type)])] -> [Doc a]
forall a b. (a -> b) -> [a] -> [b]
map (ResTag, [([DistInput], VName, Type)]) -> Doc a
forall {a} {a} {a}. (Pretty a, Pretty a) => (a, [a]) -> Doc a
onRes (ResMap -> [(ResTag, [([DistInput], VName, Type)])]
forall k a. Map k a -> [(k, a)]
M.toList ResMap
resmap) [Doc a] -> [Doc a] -> [Doc a]
forall a. Semigroup a => a -> a -> a
<> (DistRep -> Doc a) -> [DistRep] -> [Doc a]
forall a b. (a -> b) -> [a] -> [b]
map DistRep -> Doc a
forall {a} {a} {a} {ann}.
(Pretty a, Pretty a, Pretty a) =>
(a, Either a a) -> Doc ann
onRep [DistRep]
reps
      stms' :: Doc a
stms' = [Doc a] -> Doc a
forall a. [Doc a] -> Doc a
stack ([Doc a] -> Doc a) -> [Doc a] -> Doc a
forall a b. (a -> b) -> a -> b
$ (DistStm -> Doc a) -> [DistStm] -> [Doc a]
forall a b. (a -> b) -> [a] -> [b]
map DistStm -> Doc a
forall a ann. Pretty a => a -> Doc ann
forall ann. DistStm -> Doc ann
pretty ([DistStm] -> [Doc a]) -> [DistStm] -> [Doc a]
forall a b. (a -> b) -> a -> b
$ Seq DistStm -> [DistStm]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq DistStm
stms
      onRes :: (a, [a]) -> Doc a
onRes (a
rt, [a]
binds) =
        [Doc a] -> Doc a
forall a. [Doc a] -> Doc a
stack [Doc a
"let" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc a
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
v Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a
"=" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc a
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
rt | a
v <- [a]
binds]
      onRep :: (a, Either a a) -> Doc ann
onRep (a
v, Left a
se) =
        Doc ann
"let" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
v Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"=" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"rep" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
se)
      onRep (a
v, Right a
tag) =
        Doc ann
"let" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
v Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"=" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"rep" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
tag)

resultMap :: [(VName, DistInput)] -> DistStms -> Pat Type -> Result -> ResMap
resultMap :: DistInputs -> Seq DistStm -> Pat Type -> Result -> ResMap
resultMap DistInputs
avail_inputs Seq DistStm
stms Pat Type
pat Result
res = (DistStm -> ResMap) -> Seq DistStm -> ResMap
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((DistResult -> ResMap) -> [DistResult] -> ResMap
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DistResult -> ResMap
f ([DistResult] -> ResMap)
-> (DistStm -> [DistResult]) -> DistStm -> ResMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistStm -> [DistResult]
distStmResult) Seq DistStm
stms
  where
    pes :: Map VName (PatElem Type)
pes = [(VName, PatElem Type)] -> Map VName (PatElem Type)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, PatElem Type)] -> Map VName (PatElem Type))
-> [(VName, PatElem Type)] -> Map VName (PatElem Type)
forall a b. (a -> b) -> a -> b
$ do
      stm <- Seq DistStm -> [DistStm]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq DistStm
stms
      pe <- concatMap (patElems . stmPat) (distStmStms stm)
      pure (patElemName pe, pe)
    f :: DistResult -> ResMap
f (DistResult ResTag
rt DistType
_ VName
v) =
      case [([DistInput], VName, Type)]
-> (PatElem Type -> [([DistInput], VName, Type)])
-> Maybe (PatElem Type)
-> [([DistInput], VName, Type)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] PatElem Type -> [([DistInput], VName, Type)]
forall {c}. PatElem c -> [([DistInput], VName, c)]
findRess (Maybe (PatElem Type) -> [([DistInput], VName, Type)])
-> Maybe (PatElem Type) -> [([DistInput], VName, Type)]
forall a b. (a -> b) -> a -> b
$ VName -> Map VName (PatElem Type) -> Maybe (PatElem Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName (PatElem Type)
pes of
        [] -> ResMap
forall a. Monoid a => a
mempty
        [([DistInput], VName, Type)]
binds -> ResTag -> [([DistInput], VName, Type)] -> ResMap
forall k a. k -> a -> Map k a
M.singleton ResTag
rt [([DistInput], VName, Type)]
binds
    findRess :: PatElem c -> [([DistInput], VName, c)]
findRess (PatElem VName
v c
v_t) = do
      (SubExpRes cs se, pv) <- Result -> [VName] -> [(SubExpRes, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
res (Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pat)
      if se == Var v
        then pure (map findCert (unCerts cs), pv, v_t)
        else []
    findCert :: VName -> DistInput
findCert VName
v = DistInput -> Maybe DistInput -> DistInput
forall a. a -> Maybe a -> a
fromMaybe (VName -> Type -> DistInput
DistInputFree VName
v (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Unit)) (Maybe DistInput -> DistInput) -> Maybe DistInput -> DistInput
forall a b. (a -> b) -> a -> b
$ VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
avail_inputs

splitIrregDims :: Names -> Type -> (Rank, Type)
splitIrregDims :: Names -> Type -> (Rank, Type)
splitIrregDims Names
bound_outside (Array PrimType
pt Shape
shape NoMode
o) =
  let (Segments
reg, Segments
irreg) =
        (Segments -> Segments)
-> (Segments, Segments) -> (Segments, Segments)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Segments -> Segments
forall a. [a] -> [a]
reverse ((Segments, Segments) -> (Segments, Segments))
-> (Segments, Segments) -> (Segments, Segments)
forall a b. (a -> b) -> a -> b
$ (SubExp -> Bool) -> Segments -> (Segments, Segments)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span SubExp -> Bool
regDim (Segments -> (Segments, Segments))
-> Segments -> (Segments, Segments)
forall a b. (a -> b) -> a -> b
$ Segments -> Segments
forall a. [a] -> [a]
reverse (Segments -> Segments) -> Segments -> Segments
forall a b. (a -> b) -> a -> b
$ Shape -> Segments
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
   in (Int -> Rank
Rank (Int -> Rank) -> Int -> Rank
forall a b. (a -> b) -> a -> b
$ Segments -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Segments
irreg, PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt (Segments -> Shape
forall d. [d] -> ShapeBase d
Shape Segments
reg) NoMode
o)
  where
    regDim :: SubExp -> Bool
regDim (Var VName
v) = VName
v VName -> Names -> Bool
`nameIn` Names
bound_outside
    regDim Constant {} = Bool
True
splitIrregDims Names
_ Type
t = (Rank
forall a. Monoid a => a
mempty, Type
t)

freeInput :: [(VName, DistInput)] -> VName -> Maybe (VName, DistInput)
freeInput :: DistInputs -> VName -> Maybe (VName, DistInput)
freeInput DistInputs
avail_inputs VName
v =
  (VName
v,) (DistInput -> (VName, DistInput))
-> Maybe DistInput -> Maybe (VName, DistInput)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
avail_inputs

patInput :: ResTag -> PatElem Type -> (VName, DistInput)
patInput :: ResTag -> PatElem Type -> (VName, DistInput)
patInput ResTag
tag PatElem Type
pe =
  (PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe, ResTag -> Type -> DistInput
DistInput ResTag
tag (Type -> DistInput) -> Type -> DistInput
forall a b. (a -> b) -> a -> b
$ PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem Type
pe)

nextResTag :: DistInputs -> ResTag
nextResTag :: DistInputs -> ResTag
nextResTag = (ResTag -> (VName, DistInput) -> ResTag)
-> ResTag -> DistInputs -> ResTag
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ResTag -> (VName, DistInput) -> ResTag
forall {a}. ResTag -> (a, DistInput) -> ResTag
step (Int -> ResTag
ResTag Int
0)
  where
    step :: ResTag -> (a, DistInput) -> ResTag
step ResTag
next (a
_, DistInputFree VName
_ Type
_) =
      ResTag
next
    step ResTag
next (a
_, DistInput (ResTag Int
i) Type
_) =
      ResTag -> ResTag -> ResTag
forall a. Ord a => a -> a -> a
max ResTag
next (Int -> ResTag
ResTag (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))

distributeBody ::
  DistIrregularity ->
  FunHasParallelism ->
  Scope rep ->
  Segments ->
  DistInputs ->
  Body SOACS ->
  (DistInputs, DistStms)
distributeBody :: forall rep.
DistIrregularity
-> FunHasParallelism
-> Scope rep
-> Segments
-> DistInputs
-> Body SOACS
-> (DistInputs, Seq DistStm)
distributeBody DistIrregularity
irreg_mode FunHasParallelism
funHasParallelism Scope rep
outer_scope Segments
w DistInputs
param_inputs Body SOACS
body = do
  let ((ResTag
_, DistInputs
avail_inputs), Seq DistStm
stms) =
        ((ResTag, DistInputs)
 -> Stm SOACS -> ((ResTag, DistInputs), DistStm))
-> (ResTag, DistInputs)
-> Stms SOACS
-> ((ResTag, DistInputs), Seq DistStm)
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL (ResTag, DistInputs)
-> Stm SOACS -> ((ResTag, DistInputs), DistStm)
distributeStm (DistInputs -> ResTag
nextResTag DistInputs
param_inputs, DistInputs
param_inputs) (Stms SOACS -> ((ResTag, DistInputs), Seq DistStm))
-> Stms SOACS -> ((ResTag, DistInputs), Seq DistStm)
forall a b. (a -> b) -> a -> b
$
          Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body
   in ( DistInputs
avail_inputs,
        DistIrregularity
-> FunHasParallelism -> Result -> Seq DistStm -> Seq DistStm
classifyStms DistIrregularity
irreg_mode FunHasParallelism
funHasParallelism (Body SOACS -> Result
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
body) Seq DistStm
stms
      )
  where
    bound_outside :: Names
bound_outside = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Scope rep -> [VName]
forall k a. Map k a -> [k]
M.keys Scope rep
outer_scope
    distType :: Type -> DistType
distType Type
t = (Rank -> Type -> DistType) -> (Rank, Type) -> DistType
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Segments -> Rank -> Type -> DistType
DistType Segments
w) ((Rank, Type) -> DistType) -> (Rank, Type) -> DistType
forall a b. (a -> b) -> a -> b
$ Names -> Type -> (Rank, Type)
splitIrregDims Names
bound_outside Type
t
    distributeStm :: (ResTag, DistInputs)
-> Stm SOACS -> ((ResTag, DistInputs), DistStm)
distributeStm (ResTag Int
tag, DistInputs
avail_inputs) Stm SOACS
stm =
      let pat :: Pat (LetDec SOACS)
pat = Stm SOACS -> Pat (LetDec SOACS)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm SOACS
stm
          new_tags :: [ResTag]
new_tags = (Int -> ResTag) -> [Int] -> [ResTag]
forall a b. (a -> b) -> [a] -> [b]
map Int -> ResTag
ResTag ([Int] -> [ResTag]) -> [Int] -> [ResTag]
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (Pat (LetDec SOACS) -> Int
forall dec. Pat dec -> Int
patSize Pat (LetDec SOACS)
pat) [Int
tag ..]
          avail_inputs' :: DistInputs
avail_inputs' =
            DistInputs
avail_inputs DistInputs -> DistInputs -> DistInputs
forall a. Semigroup a => a -> a -> a
<> (ResTag -> PatElem Type -> (VName, DistInput))
-> [ResTag] -> [PatElem Type] -> DistInputs
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ResTag -> PatElem Type -> (VName, DistInput)
patInput [ResTag]
new_tags (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat)
          free_in_stm :: Names
free_in_stm = Stm SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Stm SOACS
stm
          used_free :: DistInputs
used_free = (VName -> Maybe (VName, DistInput)) -> [VName] -> DistInputs
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DistInputs -> VName -> Maybe (VName, DistInput)
freeInput DistInputs
avail_inputs) ([VName] -> DistInputs) -> [VName] -> DistInputs
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
free_in_stm
          used_free_types :: DistInputs
used_free_types =
            (VName -> Maybe (VName, DistInput)) -> [VName] -> DistInputs
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DistInputs -> VName -> Maybe (VName, DistInput)
freeInput DistInputs
avail_inputs)
              ([VName] -> DistInputs)
-> (DistInputs -> [VName]) -> DistInputs -> DistInputs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList
              (Names -> [VName])
-> (DistInputs -> Names) -> DistInputs -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, DistInput) -> Names) -> DistInputs -> Names
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Type -> Names
forall a. FreeIn a => a -> Names
freeIn (Type -> Names)
-> ((VName, DistInput) -> Type) -> (VName, DistInput) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistInput -> Type
distInputType (DistInput -> Type)
-> ((VName, DistInput) -> DistInput) -> (VName, DistInput) -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, DistInput) -> DistInput
forall a b. (a, b) -> b
snd)
              (DistInputs -> DistInputs) -> DistInputs -> DistInputs
forall a b. (a -> b) -> a -> b
$ DistInputs
used_free
          stm' :: DistStm
stm' =
            DistInputs -> [DistResult] -> DistBody -> DistStm
DistStm
              (DistInputs -> DistInputs
nubInputs (DistInputs -> DistInputs) -> DistInputs -> DistInputs
forall a b. (a -> b) -> a -> b
$ DistInputs
used_free_types DistInputs -> DistInputs -> DistInputs
forall a. Semigroup a => a -> a -> a
<> DistInputs
used_free)
              ((ResTag -> DistType -> VName -> DistResult)
-> [ResTag] -> [DistType] -> [VName] -> [DistResult]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 ResTag -> DistType -> VName -> DistResult
DistResult [ResTag]
new_tags ((Type -> DistType) -> [Type] -> [DistType]
forall a b. (a -> b) -> [a] -> [b]
map Type -> DistType
distType ([Type] -> [DistType]) -> [Type] -> [DistType]
forall a b. (a -> b) -> a -> b
$ Pat (LetDec SOACS) -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat (LetDec SOACS)
pat) (Pat (LetDec SOACS) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (LetDec SOACS)
pat))
              (Stm SOACS -> DistBody
ParallelStm Stm SOACS
stm)
       in ((Int -> ResTag
ResTag (Int -> ResTag) -> Int -> ResTag
forall a b. (a -> b) -> a -> b
$ Int
tag Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [ResTag] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ResTag]
new_tags, DistInputs
avail_inputs'), DistStm
stm')

isParallelDistStm :: (Stm SOACS -> Bool) -> DistStm -> Bool
isParallelDistStm :: (Stm SOACS -> Bool) -> DistStm -> Bool
isParallelDistStm Stm SOACS -> Bool
stm_is_parallel (DistStm DistInputs
_ [DistResult]
res (ParallelStm Stm SOACS
stm)) =
  Stm SOACS -> Bool
stm_is_parallel Stm SOACS
stm Bool -> Bool -> Bool
|| Bool -> Bool
not ((DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
res)
isParallelDistStm Stm SOACS -> Bool
_ DistStm
_ = Bool
False

noSequentialAttr :: Stm SOACS -> Bool
noSequentialAttr :: Stm SOACS -> Bool
noSequentialAttr Stm SOACS
stm =
  Bool -> Bool
not (Attr
"sequential" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs (Stm SOACS -> StmAux (ExpDec SOACS)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm SOACS
stm))

-- | Does the statement contain meaningful parallelism - a SOAC or a call to
-- a parallel function, possibly nested inside sequential control flow?
-- Basic operations such as 'Iota' or 'Replicate' do not count. See Note
-- [Meaningful Parallelism].
stmHasMeaningfulParallelism :: FunHasParallelism -> Stm SOACS -> Bool
stmHasMeaningfulParallelism :: FunHasParallelism -> Stm SOACS -> Bool
stmHasMeaningfulParallelism FunHasParallelism
funHasParallelism = Stm SOACS -> Bool
hasParallelism
  where
    hasParallelism :: Stm SOACS -> Bool
hasParallelism Stm SOACS
stm' =
      Stm SOACS -> Bool
noSequentialAttr Stm SOACS
stm'
        Bool -> Bool -> Bool
&& case Stm SOACS -> Exp SOACS
forall rep. Stm rep -> Exp rep
stmExp Stm SOACS
stm' of
          BasicOp BasicOp
_ -> Bool
False
          Apply Name
fname [(SubExp, Diet)]
_ [(RetType SOACS, RetAls)]
_ Safety
_ -> FunHasParallelism
funHasParallelism Name
fname
          Match Segments
_ [Case (Body SOACS)]
cases Body SOACS
def_case MatchDec (BranchType SOACS)
_ ->
            (Stm SOACS -> Bool) -> Stms SOACS -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Stm SOACS -> Bool
hasParallelism (Stms SOACS -> Bool) -> Stms SOACS -> Bool
forall a b. (a -> b) -> a -> b
$
              Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
def_case
                Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> [Stms SOACS] -> Stms SOACS
forall a. Monoid a => [a] -> a
mconcat ((Case (Body SOACS) -> Stms SOACS)
-> [Case (Body SOACS)] -> [Stms SOACS]
forall a b. (a -> b) -> [a] -> [b]
map (Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS)
-> (Case (Body SOACS) -> Body SOACS)
-> Case (Body SOACS)
-> Stms SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (Body SOACS) -> Body SOACS
forall body. Case body -> body
caseBody) [Case (Body SOACS)]
cases)
          Loop [(FParam SOACS, SubExp)]
_ LoopForm
_ Body SOACS
body -> (Stm SOACS -> Bool) -> Stms SOACS -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Stm SOACS -> Bool
hasParallelism (Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body)
          WithAcc [WithAccInput SOACS]
_ Lambda SOACS
lam -> (Stm SOACS -> Bool) -> Stms SOACS -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Stm SOACS -> Bool
hasParallelism (Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam))
          Op Op SOACS
op -> SOAC SOACS -> Bool
forall {rep}. SOAC rep -> Bool
isParallelOp Op SOACS
SOAC SOACS
op

    isParallelOp :: SOAC rep -> Bool
isParallelOp JVP {} = String -> Bool
forall a. HasCallStack => String -> a
error String
"stmHasMeaningfulParallelism: JVP"
    isParallelOp VJP {} = String -> Bool
forall a. HasCallStack => String -> a
error String
"stmHasMeaningfulParallelism: VJP"
    isParallelOp SOAC rep
_ = Bool
True

isParallelStm :: FunHasParallelism -> Stm SOACS -> Bool
isParallelStm :: FunHasParallelism -> Stm SOACS -> Bool
isParallelStm FunHasParallelism
funHasParallelism Stm SOACS
stm =
  Stm SOACS -> Bool
noSequentialAttr Stm SOACS
stm
    Bool -> Bool -> Bool
&& (Exp SOACS -> Bool
forall {rep}. Exp rep -> Bool
parallelBasicOp (Stm SOACS -> Exp SOACS
forall rep. Stm rep -> Exp rep
stmExp Stm SOACS
stm) Bool -> Bool -> Bool
|| FunHasParallelism -> Stm SOACS -> Bool
stmHasMeaningfulParallelism FunHasParallelism
funHasParallelism Stm SOACS
stm)
  where
    parallelBasicOp :: Exp rep -> Bool
parallelBasicOp (BasicOp BasicOp
op) = BasicOp -> Bool
isParallelBasicOp BasicOp
op
    parallelBasicOp Exp rep
_ = Bool
False

    isParallelBasicOp :: BasicOp -> Bool
isParallelBasicOp (Update Safety
_ VName
_ Slice SubExp
slice SubExp
_) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Segments -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Segments -> Bool) -> Segments -> Bool
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> Segments
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
    isParallelBasicOp Concat {} = Bool
True
    isParallelBasicOp Iota {} = Bool
True
    isParallelBasicOp Replicate {} = Bool
True
    isParallelBasicOp (FlatUpdate VName
_ FlatSlice SubExp
flat_slice VName
_) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Segments -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Segments -> Bool) -> Segments -> Bool
forall a b. (a -> b) -> a -> b
$ FlatSlice SubExp -> Segments
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice
    isParallelBasicOp Manifest {} = Bool
True
    isParallelBasicOp Rearrange {} = Bool
True
    isParallelBasicOp Reshape {} = Bool
True
    isParallelBasicOp (FlatIndex VName
_ FlatSlice SubExp
flat_slice) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Segments -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Segments -> Bool) -> Segments -> Bool
forall a b. (a -> b) -> a -> b
$ FlatSlice SubExp -> Segments
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice
    isParallelBasicOp (Index VName
_ Slice SubExp
slice) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Segments -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Segments -> Bool) -> Segments -> Bool
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> Segments
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
    -- The work done by an 'ArrayLit' is its number of rows times the size of a
    -- row. Since the size is syntactic, we consider a literal of scalars to be
    -- always cheap enough for a single thread to construct.
    isParallelBasicOp (ArrayLit Segments
_ Type
row_t) = Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
row_t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    -- Now the sequential ones - we handle them explicitly so we will notice if
    -- we ever add a new one.
    isParallelBasicOp ArrayVal {} = Bool
False
    isParallelBasicOp Scratch {} = Bool
False
    isParallelBasicOp SubExp {} = Bool
False
    isParallelBasicOp Opaque {} = Bool
False
    isParallelBasicOp UnOp {} = Bool
False
    isParallelBasicOp BinOp {} = Bool
False
    isParallelBasicOp CmpOp {} = Bool
False
    isParallelBasicOp ConvOp {} = Bool
False
    isParallelBasicOp Assert {} = Bool
False
    isParallelBasicOp UpdateAcc {} = Bool
False
    isParallelBasicOp UserParam {} = Bool
False

isRegularDistResult :: DistResult -> Bool
isRegularDistResult :: DistResult -> Bool
isRegularDistResult (DistResult ResTag
_ (DistType Segments
_ (Rank Int
r) Type
_) VName
_) = Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0

-- | Does the statement contain, inside a 'Loop' or 'Match', a statement whose
-- result is sized by a name bound within the enclosing statement itself? Such a
-- statement can never be executed sequentially inside a GPU kernel, as it
-- implies an allocation whose size cannot be computed before the kernel is
-- launched; only the machinery for flattening irregular arrays can handle it.
-- In contrast, sizes that are nonuniform merely by being variant to the
-- enclosing map-nest are fine, as memory expansion can compute those via
-- slicing. This is essentially a heuristic where we bet that slicing is
-- efficient; the fully principled stance would be to not allow any
-- nonuniformity. This is one of the criteria of 'mustDistribute' in
-- 'classifyStms'. See Note [Meaningful Parallelism].
stmHasNonuniformInside :: Stm SOACS -> Bool
stmHasNonuniformInside :: Stm SOACS -> Bool
stmHasNonuniformInside = Names -> Exp SOACS -> Bool
forall {rep}. Typed (LetDec rep) => Names -> Exp rep -> Bool
inExp Names
forall a. Monoid a => a
mempty (Exp SOACS -> Bool)
-> (Stm SOACS -> Exp SOACS) -> Stm SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm SOACS -> Exp SOACS
forall rep. Stm rep -> Exp rep
stmExp
  where
    nonuniform :: Names -> TypeBase Shape m -> Bool
nonuniform Names
bound = (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
bound) ([VName] -> Bool)
-> (TypeBase Shape m -> [VName]) -> TypeBase Shape m -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Segments -> [VName]
subExpVars (Segments -> [VName])
-> (TypeBase Shape m -> Segments) -> TypeBase Shape m -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape m -> Segments
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims
    inExp :: Names -> Exp rep -> Bool
inExp Names
bound Exp rep
e =
      case Exp rep
e of
        Loop [(FParam rep, SubExp)]
merge LoopForm
form Body rep
body ->
          Names -> Body rep -> Bool
inBody
            ( Names
bound
                Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList (((FParam rep, SubExp) -> VName)
-> [(FParam rep, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (FParam rep -> VName
forall dec. Param dec -> VName
paramName (FParam rep -> VName)
-> ((FParam rep, SubExp) -> FParam rep)
-> (FParam rep, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> FParam rep
forall a b. (a, b) -> a
fst) [(FParam rep, SubExp)]
merge)
                Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList (Map VName (NameInfo (ZonkAny 0)) -> [VName]
forall k a. Map k a -> [k]
M.keys (LoopForm -> Map VName (NameInfo (ZonkAny 0))
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form))
            )
            Body rep
body
        Match Segments
_ [Case (Body rep)]
cases Body rep
def_body MatchDec (BranchType rep)
_ ->
          (Body rep -> Bool) -> [Body rep] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Names -> Body rep -> Bool
inBody Names
bound) (Body rep
def_body Body rep -> [Body rep] -> [Body rep]
forall a. a -> [a] -> [a]
: (Case (Body rep) -> Body rep) -> [Case (Body rep)] -> [Body rep]
forall a b. (a -> b) -> [a] -> [b]
map Case (Body rep) -> Body rep
forall body. Case body -> body
caseBody [Case (Body rep)]
cases)
        WithAcc [WithAccInput rep]
_ Lambda rep
lam -> Names -> Body rep -> Bool
inBody Names
bound (Lambda rep -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam)
        Exp rep
_ -> Bool
False
    inBody :: Names -> Body rep -> Bool
inBody Names
bound Body rep
body =
      (Stm rep -> Bool) -> Seq (Stm rep) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Names -> Stm rep -> Bool
inStm (Names
bound Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Body rep -> Names
forall rep. Body rep -> Names
boundInBody Body rep
body)) (Seq (Stm rep) -> Bool) -> Seq (Stm rep) -> Bool
forall a b. (a -> b) -> a -> b
$ Body rep -> Seq (Stm rep)
forall rep res. GBody rep res -> Stms rep
bodyStms Body rep
body
    inStm :: Names -> Stm rep -> Bool
inStm Names
bound Stm rep
stm =
      (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Names -> Type -> Bool
forall {m}. Names -> TypeBase Shape m -> Bool
nonuniform Names
bound) (Pat (LetDec rep) -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes (Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm))
        Bool -> Bool -> Bool
|| Names -> Exp rep -> Bool
inExp Names
bound (Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm)

-- | Merge a group of scalar 'DistStm's into a single one.
mergeGroup :: Result -> DistStms -> DistStms -> DistStm
mergeGroup :: Result -> Seq DistStm -> Seq DistStm -> DistStm
mergeGroup Result
bodyRes Seq DistStm
ds Seq DistStm
rest =
  let resTags :: Set ResTag
resTags =
        [ResTag] -> Set ResTag
forall a. Ord a => [a] -> Set a
S.fromList ([ResTag] -> Set ResTag) -> [ResTag] -> Set ResTag
forall a b. (a -> b) -> a -> b
$ (DistStm -> [ResTag]) -> Seq DistStm -> [ResTag]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((DistResult -> ResTag) -> [DistResult] -> [ResTag]
forall a b. (a -> b) -> [a] -> [b]
map DistResult -> ResTag
distResTag ([DistResult] -> [ResTag])
-> (DistStm -> [DistResult]) -> DistStm -> [ResTag]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistStm -> [DistResult]
distStmResult) Seq DistStm
ds
      isInternal :: (a, DistInput) -> Bool
isInternal (a
_, DistInput ResTag
rt Type
_) = ResTag
rt ResTag -> Set ResTag -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set ResTag
resTags
      isInternal (a, DistInput)
_ = Bool
False
      externalInputs :: DistInputs
externalInputs =
        DistInputs -> DistInputs
nubInputs (DistInputs -> DistInputs) -> DistInputs -> DistInputs
forall a b. (a -> b) -> a -> b
$
          (DistStm -> DistInputs) -> Seq DistStm -> DistInputs
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (((VName, DistInput) -> Bool) -> DistInputs -> DistInputs
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> ((VName, DistInput) -> Bool) -> (VName, DistInput) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, DistInput) -> Bool
forall {a}. (a, DistInput) -> Bool
isInternal) (DistInputs -> DistInputs)
-> (DistStm -> DistInputs) -> DistStm -> DistInputs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistStm -> DistInputs
distStmInputs) Seq DistStm
ds
      externalResults :: [DistResult]
externalResults =
        [DistResult] -> [DistResult]
forall a. Ord a => [a] -> [a]
nubOrd ([DistResult] -> [DistResult]) -> [DistResult] -> [DistResult]
forall a b. (a -> b) -> a -> b
$
          (DistStm -> [DistResult]) -> Seq DistStm -> [DistResult]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((DistResult -> Bool) -> [DistResult] -> [DistResult]
forall a. (a -> Bool) -> [a] -> [a]
filter (Result -> Seq DistStm -> DistResult -> Bool
isExternal Result
bodyRes Seq DistStm
rest) ([DistResult] -> [DistResult])
-> (DistStm -> [DistResult]) -> DistStm -> [DistResult]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistStm -> [DistResult]
distStmResult) Seq DistStm
ds
      allStms :: Stms SOACS
allStms = (DistStm -> Stms SOACS) -> Seq DistStm -> Stms SOACS
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DistStm -> Stms SOACS
distStmStms Seq DistStm
ds
   in DistInputs -> [DistResult] -> DistBody -> DistStm
DistStm DistInputs
externalInputs [DistResult]
externalResults (Stms SOACS -> DistBody
ScalarStm Stms SOACS
allStms)

groupStms ::
  (DistStm -> Bool) ->
  Result ->
  Seq.Seq DistStm ->
  Seq.Seq DistStm
groupStms :: (DistStm -> Bool) -> Result -> Seq DistStm -> Seq DistStm
groupStms DistStm -> Bool
_ Result
_ Seq DistStm
Seq.Empty = Seq DistStm
forall a. Monoid a => a
mempty
groupStms DistStm -> Bool
dist_stm_is_parallel Result
body_res Seq DistStm
ds' =
  let (Seq DistStm
scalars, Seq DistStm
rest) = (DistStm -> Bool) -> Seq DistStm -> (Seq DistStm, Seq DistStm)
forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
Seq.breakl DistStm -> Bool
dist_stm_is_parallel Seq DistStm
ds'
      scalar_grouped :: Seq DistStm
scalar_grouped
        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Seq DistStm -> Bool
forall a. Seq a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Seq DistStm
scalars =
            DistStm -> Seq DistStm
forall a. a -> Seq a
Seq.singleton (DistStm -> Seq DistStm) -> DistStm -> Seq DistStm
forall a b. (a -> b) -> a -> b
$ Result -> Seq DistStm -> Seq DistStm -> DistStm
mergeGroup Result
body_res Seq DistStm
scalars Seq DistStm
rest
        | Bool
otherwise = Seq DistStm
forall a. Monoid a => a
mempty
   in case Seq DistStm
rest of
        Seq DistStm
Seq.Empty -> Seq DistStm
scalar_grouped
        DistStm
p Seq.:<| Seq DistStm
ps ->
          Seq DistStm
scalar_grouped Seq DistStm -> Seq DistStm -> Seq DistStm
forall a. Semigroup a => a -> a -> a
<> (DistStm
p DistStm -> Seq DistStm -> Seq DistStm
forall a. a -> Seq a -> Seq a
Seq.<| (DistStm -> Bool) -> Result -> Seq DistStm -> Seq DistStm
groupStms DistStm -> Bool
dist_stm_is_parallel Result
body_res Seq DistStm
ps)

--  we should probably sort the DistStms first and we should assume they are sorted
-- and then given to this function.
classifyStms :: DistIrregularity -> FunHasParallelism -> Result -> DistStms -> DistStms
classifyStms :: DistIrregularity
-> FunHasParallelism -> Result -> Seq DistStm -> Seq DistStm
classifyStms DistIrregularity
irreg_mode FunHasParallelism
funHasParallelism Result
body_res = Seq DistStm -> Seq DistStm
classify
  where
    -- Distribute the statements that are parallel, plus those that
    -- 'mustDistribute' regardless of parallelism. If no statement contains
    -- meaningful parallelism, no statement counts as parallel, so that the
    -- trivially parallel statements are treated as sequential as well. See Note
    -- [Meaningful Parallelism].
    --
    -- With 'SequentialiseIrregularBasicOps', nonuniform basic operations
    -- further count as parallel only when their irregular arrays escape the
    -- scalar group. With 'SequentialiseIrregularAll', this goes for any
    -- statement, not just basic operations.
    classify :: Seq DistStm -> Seq DistStm
classify Seq DistStm
ds =
      let -- Which statements are candidates for sequentialising when their
          -- nonuniformity does not escape the scalar group.
          sequentialisable :: DistStm -> Bool
sequentialisable = case DistIrregularity
irreg_mode of
            DistIrregularity
DistributeIrregular -> Bool -> DistStm -> Bool
forall a b. a -> b -> a
const Bool
False
            DistIrregularity
SequentialiseIrregularBasicOps -> DistStm -> Bool
isBasicOpDistStm
            DistIrregularity
SequentialiseIrregularAll -> Bool -> DistStm -> Bool
forall a b. a -> b -> a
const Bool
True
          parallel :: DistStm -> Bool
parallel
            | (DistStm -> Bool) -> Seq DistStm -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any DistStm -> Bool
meaningfulDistStm Seq DistStm
ds =
                \DistStm
d ->
                  (Stm SOACS -> Bool) -> DistStm -> Bool
isParallelDistStm (FunHasParallelism -> Stm SOACS -> Bool
isParallelStm FunHasParallelism
funHasParallelism) DistStm
d
                    Bool -> Bool -> Bool
&& Bool -> Bool
not (DistStm -> Bool
sequentialisable DistStm
d Bool -> Bool -> Bool
&& Seq DistStm -> DistStm -> Bool
forall {t :: * -> *}. Foldable t => t DistStm -> DistStm -> Bool
involvesIrregularity Seq DistStm
ds DistStm
d)
            | Bool
otherwise = Bool -> DistStm -> Bool
forall a b. a -> b -> a
const Bool
False
          forced :: Set DistStm
forced = Set DistStm -> Seq DistStm -> Set DistStm
forall {t :: * -> *}.
Foldable t =>
Set DistStm -> t DistStm -> Set DistStm
mustDistribute ([DistStm] -> Set DistStm
forall a. Ord a => [a] -> Set a
S.fromList ([DistStm] -> Set DistStm) -> [DistStm] -> Set DistStm
forall a b. (a -> b) -> a -> b
$ (DistStm -> Bool) -> [DistStm] -> [DistStm]
forall a. (a -> Bool) -> [a] -> [a]
filter DistStm -> Bool
parallel ([DistStm] -> [DistStm]) -> [DistStm] -> [DistStm]
forall a b. (a -> b) -> a -> b
$ Seq DistStm -> [DistStm]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq DistStm
ds) Seq DistStm
ds
       in (DistStm -> Bool) -> Result -> Seq DistStm -> Seq DistStm
groupStms (\DistStm
d -> DistStm -> Bool
parallel DistStm
d Bool -> Bool -> Bool
|| DistStm
d DistStm -> Set DistStm -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set DistStm
forced) Result
body_res Seq DistStm
ds

    meaningfulDistStm :: DistStm -> Bool
meaningfulDistStm d :: DistStm
d@(DistStm DistInputs
_ [DistResult]
_ (ParallelStm Stm SOACS
stm)) =
      (DistStm -> Bool
isBasicOpDistStm DistStm
d Bool -> Bool -> Bool
&& (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Result -> Names
forall a. FreeIn a => a -> Names
freeIn Result
body_res) (Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames (Stm SOACS -> Pat (LetDec SOACS)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm SOACS
stm)))
        Bool -> Bool -> Bool
|| FunHasParallelism -> Stm SOACS -> Bool
stmHasMeaningfulParallelism FunHasParallelism
funHasParallelism Stm SOACS
stm
    meaningfulDistStm DistStm
_ = Bool
False

    isBasicOpDistStm :: DistStm -> Bool
isBasicOpDistStm (DistStm DistInputs
_ [DistResult]
_ (ParallelStm (Let Pat (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ BasicOp {}))) = Bool
True
    isBasicOpDistStm DistStm
_ = Bool
False

    distStmHasNonuniformInside :: DistStm -> Bool
distStmHasNonuniformInside (DistStm DistInputs
_ [DistResult]
_ (ParallelStm Stm SOACS
stm)) =
      Stm SOACS -> Bool
stmHasNonuniformInside Stm SOACS
stm
    distStmHasNonuniformInside DistStm
_ = Bool
False

    -- Whether a statement produces an irregular array, or consumes
    -- one produced elsewhere in the body.
    involvesIrregularity :: t DistStm -> DistStm -> Bool
involvesIrregularity t DistStm
ds DistStm
d =
      Bool -> Bool
not ((DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult (DistStm -> [DistResult]
distStmResult DistStm
d))
        Bool -> Bool -> Bool
|| ((VName, DistInput) -> Bool) -> DistInputs -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName, DistInput) -> Bool
forall {a}. (a, DistInput) -> Bool
consumesIrregular (DistStm -> DistInputs
distStmInputs DistStm
d)
      where
        consumesIrregular :: (a, DistInput) -> Bool
consumesIrregular (a
_, DistInput ResTag
rt Type
_) = ResTag
rt ResTag -> Set ResTag -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set ResTag
irregular_tags
        consumesIrregular (a, DistInput)
_ = Bool
False
        irregular_tags :: Set ResTag
irregular_tags =
          [ResTag] -> Set ResTag
forall a. Ord a => [a] -> Set a
S.fromList ([ResTag] -> Set ResTag) -> [ResTag] -> Set ResTag
forall a b. (a -> b) -> a -> b
$
            (DistResult -> ResTag) -> [DistResult] -> [ResTag]
forall a b. (a -> b) -> [a] -> [b]
map DistResult -> ResTag
distResTag ([DistResult] -> [ResTag]) -> [DistResult] -> [ResTag]
forall a b. (a -> b) -> a -> b
$
              (DistResult -> Bool) -> [DistResult] -> [DistResult]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (DistResult -> Bool) -> DistResult -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistResult -> Bool
isRegularDistResult) ([DistResult] -> [DistResult]) -> [DistResult] -> [DistResult]
forall a b. (a -> b) -> a -> b
$
                (DistStm -> [DistResult]) -> t DistStm -> [DistResult]
forall m a. Monoid m => (a -> m) -> t a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DistStm -> [DistResult]
distStmResult t DistStm
ds

    -- The statements that require distribution regardless of whether they
    -- contain profitable parallelism, because a sequentially executed scalar
    -- group cannot handle their sizes:
    --
    -- (1) Statements whose irregular results are used outside the scalar
    --     group (by the body result, or transitively by another distributed
    --     statement), as arrays produced by sequentially executed groups
    --     must be regular.
    --
    -- (2) Compound statements (e.g. loops) with irregular results, even
    --     internally used ones: they contain allocations of nonuniform size
    --     that only flattening can handle, while a basic operation can
    --     reasonably be executed sequentially by a single thread when its
    --     result stays internal. This criterion does not apply under
    --     'SequentialiseIrregularAll', which is only imposed on user request,
    --     and may result in non-compileable code due to impossible memory
    --     expansion.
    --
    -- (3) Statements with sizes bound inside their own sequential control
    --     flow ('stmHasNonuniformInside').
    --
    -- The initial set of forced statements may be seeded with the statements
    -- already known to be distributed, such that irregular arrays they consume
    -- are also distributed.
    mustDistribute :: Set DistStm -> t DistStm -> Set DistStm
mustDistribute Set DistStm
seed t DistStm
ds = Set DistStm -> Set DistStm
fixpoint Set DistStm
seed
      where
        body_res_names :: Names
body_res_names = Result -> Names
forall a. FreeIn a => a -> Names
freeIn Result
body_res
        irregularResults :: DistStm -> [DistResult]
irregularResults DistStm
stm =
          (DistResult -> Bool) -> [DistResult] -> [DistResult]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (DistResult -> Bool) -> DistResult -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistResult -> Bool
isRegularDistResult) ([DistResult] -> [DistResult]) -> [DistResult] -> [DistResult]
forall a b. (a -> b) -> a -> b
$ DistStm -> [DistResult]
distStmResult DistStm
stm
        forcedTags :: Set DistStm -> Set ResTag
forcedTags Set DistStm
forced =
          [ResTag] -> Set ResTag
forall a. Ord a => [a] -> Set a
S.fromList
            [ResTag
rt | DistStm
stm <- Set DistStm -> [DistStm]
forall a. Set a -> [a]
S.toList Set DistStm
forced, (VName
_, DistInput ResTag
rt Type
_) <- DistStm -> DistInputs
distStmInputs DistStm
stm]
        compoundIrregular :: DistStm -> Bool
compoundIrregular DistStm
stm =
          DistIrregularity
irreg_mode DistIrregularity -> DistIrregularity -> Bool
forall a. Eq a => a -> a -> Bool
/= DistIrregularity
SequentialiseIrregularAll
            Bool -> Bool -> Bool
&& Bool -> Bool
not (DistStm -> Bool
isBasicOpDistStm DistStm
stm)
            Bool -> Bool -> Bool
&& Bool -> Bool
not ([DistResult] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (DistStm -> [DistResult]
irregularResults DistStm
stm))
        isForced :: Set DistStm -> DistStm -> Bool
isForced Set DistStm
forced DistStm
stm =
          DistStm -> Bool
compoundIrregular DistStm
stm
            Bool -> Bool -> Bool
|| DistStm -> Bool
distStmHasNonuniformInside DistStm
stm
            Bool -> Bool -> Bool
|| (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
              ( \DistResult
r ->
                  DistResult -> VName
distResName DistResult
r VName -> Names -> Bool
`nameIn` Names
body_res_names
                    Bool -> Bool -> Bool
|| DistResult -> ResTag
distResTag DistResult
r ResTag -> Set ResTag -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set DistStm -> Set ResTag
forcedTags Set DistStm
forced
              )
              (DistStm -> [DistResult]
irregularResults DistStm
stm)
        fixpoint :: Set DistStm -> Set DistStm
fixpoint Set DistStm
forced =
          let forced' :: Set DistStm
forced' = Set DistStm
seed Set DistStm -> Set DistStm -> Set DistStm
forall a. Semigroup a => a -> a -> a
<> [DistStm] -> Set DistStm
forall a. Ord a => [a] -> Set a
S.fromList ((DistStm -> Bool) -> [DistStm] -> [DistStm]
forall a. (a -> Bool) -> [a] -> [a]
filter (Set DistStm -> DistStm -> Bool
isForced Set DistStm
forced) ([DistStm] -> [DistStm]) -> [DistStm] -> [DistStm]
forall a b. (a -> b) -> a -> b
$ t DistStm -> [DistStm]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t DistStm
ds)
           in if Set DistStm
forced' Set DistStm -> Set DistStm -> Bool
forall a. Eq a => a -> a -> Bool
== Set DistStm
forced then Set DistStm
forced else Set DistStm -> Set DistStm
fixpoint Set DistStm
forced'

-- | A result is external if it is used by a subsequent 'DistStm' or by the body
-- result.
isExternal :: Result -> DistStms -> DistResult -> Bool
isExternal :: Result -> Seq DistStm -> DistResult -> Bool
isExternal Result
bodyRes Seq DistStm
rest (DistResult ResTag
rt DistType
_ VName
rn) =
  ResTag
rt ResTag -> Set ResTag -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set ResTag
usedByRest Bool -> Bool -> Bool
|| VName
rn VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bodyResVars Bool -> Bool -> Bool
|| VName
rn VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bodyResCerts
  where
    usedByRest :: Set ResTag
usedByRest =
      [ResTag] -> Set ResTag
forall a. Ord a => [a] -> Set a
S.fromList
        [ResTag
rt' | (VName
_, DistInput ResTag
rt' Type
_) <- (DistStm -> DistInputs) -> Seq DistStm -> DistInputs
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DistStm -> DistInputs
distStmInputs Seq DistStm
rest]
    bodyResVars :: Set VName
bodyResVars =
      [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$
        (SubExpRes -> Maybe VName) -> Result -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
          ( \(SubExpRes Certs
_ SubExp
se) -> case SubExp
se of
              Var VName
v -> VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
              SubExp
_ -> Maybe VName
forall a. Maybe a
Nothing
          )
          Result
bodyRes
    bodyResCerts :: Set VName
bodyResCerts =
      [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$
        (SubExpRes -> [VName]) -> Result -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(SubExpRes Certs
cs SubExp
_) -> Certs -> [VName]
unCerts Certs
cs) Result
bodyRes

-- | The input we are mapping over in 'distributeMap'.
data MapArray t
  = -- | A straightforward array passed in to a top-level map.
    MapArray VName Type
  | -- | Something more exotic - distribution will assign it a 'ResTag', but not
    -- do anything else. This is used to distributed nested maps whose inputs
    -- are produced in the outer nests.
    MapOther t Type

mapArrayRowType :: MapArray t -> Type
mapArrayRowType :: forall t. MapArray t -> Type
mapArrayRowType (MapArray VName
_ Type
t) = Type
t
mapArrayRowType (MapOther t
_ Type
t) = Type
t

-- This is used to handle those results that are constants or lambda
-- parameters.
findReps :: [(VName, DistInput)] -> Pat Type -> Lambda SOACS -> [DistRep]
findReps :: DistInputs -> Pat Type -> Lambda SOACS -> [DistRep]
findReps DistInputs
avail_inputs Pat Type
map_pat Lambda SOACS
lam =
  ((PatElem Type, SubExpRes) -> Maybe DistRep)
-> [(PatElem Type, SubExpRes)] -> [DistRep]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (PatElem Type, SubExpRes) -> Maybe DistRep
forall {dec}. (PatElem dec, SubExpRes) -> Maybe DistRep
f ([(PatElem Type, SubExpRes)] -> [DistRep])
-> [(PatElem Type, SubExpRes)] -> [DistRep]
forall a b. (a -> b) -> a -> b
$ [PatElem Type] -> Result -> [(PatElem Type, SubExpRes)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
map_pat) (Body SOACS -> Result
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam))
  where
    f :: (PatElem dec, SubExpRes) -> Maybe DistRep
f (PatElem dec
pe, SubExpRes Certs
_ (Var VName
v)) =
      case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
avail_inputs of
        Maybe DistInput
Nothing -> DistRep -> Maybe DistRep
forall a. a -> Maybe a
Just (PatElem dec -> VName
forall dec. PatElem dec -> VName
patElemName PatElem dec
pe, SubExp -> Either SubExp DistInput
forall a b. a -> Either a b
Left (SubExp -> Either SubExp DistInput)
-> SubExp -> Either SubExp DistInput
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v)
        Just DistInput
inp
          | VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam) ->
              DistRep -> Maybe DistRep
forall a. a -> Maybe a
Just (PatElem dec -> VName
forall dec. PatElem dec -> VName
patElemName PatElem dec
pe, DistInput -> Either SubExp DistInput
forall a b. b -> Either a b
Right DistInput
inp)
          | Bool
otherwise -> Maybe DistRep
forall a. Maybe a
Nothing
    f (PatElem dec
pe, SubExpRes Certs
_ (Constant PrimValue
v)) = do
      DistRep -> Maybe DistRep
forall a. a -> Maybe a
Just (PatElem dec -> VName
forall dec. PatElem dec -> VName
patElemName PatElem dec
pe, SubExp -> Either SubExp DistInput
forall a b. a -> Either a b
Left (SubExp -> Either SubExp DistInput)
-> SubExp -> Either SubExp DistInput
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
v)

distributeMap ::
  DistIrregularity ->
  FunHasParallelism ->
  Scope rep ->
  Pat Type ->
  Segments ->
  [MapArray t] ->
  Lambda SOACS ->
  (Distributed, M.Map ResTag t)
distributeMap :: forall rep t.
DistIrregularity
-> FunHasParallelism
-> Scope rep
-> Pat Type
-> Segments
-> [MapArray t]
-> Lambda SOACS
-> (Distributed, Map ResTag t)
distributeMap DistIrregularity
irreg_mode FunHasParallelism
funHasParallelism Scope rep
outer_scope Pat Type
map_pat Segments
w [MapArray t]
arrs Lambda SOACS
lam =
  let ((ResTag
_, Map ResTag t
arrmap), DistInputs
param_inputs) =
        ((ResTag, Map ResTag t)
 -> (Param Type, MapArray t)
 -> ((ResTag, Map ResTag t), (VName, DistInput)))
-> (ResTag, Map ResTag t)
-> [(Param Type, MapArray t)]
-> ((ResTag, Map ResTag t), DistInputs)
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL (ResTag, Map ResTag t)
-> (Param Type, MapArray t)
-> ((ResTag, Map ResTag t), (VName, DistInput))
forall {dec} {a}.
Typed dec =>
(ResTag, Map ResTag a)
-> (Param dec, MapArray a)
-> ((ResTag, Map ResTag a), (VName, DistInput))
paramInput (Int -> ResTag
ResTag Int
0, Map ResTag t
forall a. Monoid a => a
mempty) ([(Param Type, MapArray t)]
 -> ((ResTag, Map ResTag t), DistInputs))
-> [(Param Type, MapArray t)]
-> ((ResTag, Map ResTag t), DistInputs)
forall a b. (a -> b) -> a -> b
$
          [Param Type] -> [MapArray t] -> [(Param Type, MapArray t)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam) [MapArray t]
arrs
      (DistInputs
avail_inputs, Seq DistStm
stms) =
        DistIrregularity
-> FunHasParallelism
-> Scope rep
-> Segments
-> DistInputs
-> Body SOACS
-> (DistInputs, Seq DistStm)
forall rep.
DistIrregularity
-> FunHasParallelism
-> Scope rep
-> Segments
-> DistInputs
-> Body SOACS
-> (DistInputs, Seq DistStm)
distributeBody DistIrregularity
irreg_mode FunHasParallelism
funHasParallelism Scope rep
outer_scope Segments
w DistInputs
param_inputs (Body SOACS -> (DistInputs, Seq DistStm))
-> Body SOACS -> (DistInputs, Seq DistStm)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
      resmap :: ResMap
resmap =
        DistInputs -> Seq DistStm -> Pat Type -> Result -> ResMap
resultMap DistInputs
avail_inputs Seq DistStm
stms Pat Type
map_pat (Result -> ResMap) -> Result -> ResMap
forall a b. (a -> b) -> a -> b
$
          Body SOACS -> Result
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam)
      reps :: [DistRep]
reps = DistInputs -> Pat Type -> Lambda SOACS -> [DistRep]
findReps DistInputs
avail_inputs Pat Type
map_pat Lambda SOACS
lam
   in ( Seq DistStm -> DistResults -> Distributed
Distributed Seq DistStm
stms (DistResults -> Distributed) -> DistResults -> Distributed
forall a b. (a -> b) -> a -> b
$ ResMap -> [DistRep] -> DistResults
DistResults ResMap
resmap [DistRep]
reps,
        Map ResTag t
arrmap
      )
  where
    paramInput :: (ResTag, Map ResTag a)
-> (Param dec, MapArray a)
-> ((ResTag, Map ResTag a), (VName, DistInput))
paramInput (ResTag Int
i, Map ResTag a
m) (Param dec
p, MapArray VName
arr Type
_) =
      ( (Int -> ResTag
ResTag Int
i, Map ResTag a
m),
        (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
p, VName -> Type -> DistInput
DistInputFree VName
arr (Type -> DistInput) -> Type -> DistInput
forall a b. (a -> b) -> a -> b
$ Param dec -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param dec
p)
      )
    paramInput (ResTag Int
i, Map ResTag a
m) (Param dec
p, MapOther a
x Type
_) =
      ( (Int -> ResTag
ResTag (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1), ResTag -> a -> Map ResTag a -> Map ResTag a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (Int -> ResTag
ResTag Int
i) a
x Map ResTag a
m),
        (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
p, ResTag -> Type -> DistInput
DistInput (Int -> ResTag
ResTag Int
i) (Type -> DistInput) -> Type -> DistInput
forall a b. (a -> b) -> a -> b
$ Param dec -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param dec
p)
      )

-- Note [Meaningful Parallelism]
--
-- Many basic operations (Iota, Replicate, Concat, nontrivial slicing, and so
-- on) are parallel in principle, but distributing one on its own gains us
-- nothing: it performs the same work as a sequential per-thread traversal, at
-- the cost of manifesting intermediate arrays and launching extra kernels. Our
-- rule of thumb is that code contains *meaningful* parallelism only when it
-- contains at least a Screma, Hist, or a call to a parallel function - possibly
-- nested inside sequential control flow. This is what
-- 'stmHasMeaningfulParallelism' checks. It is intended to avoid "parallelising"
-- map bodies that are essentially sequential loops, but happen to use some
-- basic operations to construct arrays they operate on. This is a somewhat
-- crude classification, and it is possible to imagine a more sophisticated one
-- based on a cost-model. Note that incremental flattening (and auto-tuning)
-- might allow us to generate and pick the sequentialised versions anyway, but
-- it is more efficient to not generate them in the first place.
--
-- The notion is used in two places:
--
-- (1) When classifying the statements of a distributed body as parallel or
--     sequential ('classifyStms'): if no statement contains meaningful
--     parallelism, then the trivially parallel statements are also grouped
--     as sequential, so that ideally the entire body becomes a single
--     segmented operation.
--
-- (2) When deciding whether to multi-version a Screma
--     ('factorScremaForParallelism' in Futhark.Pass.Flatten.Incremental): a
--     fully flattened code version is worth generating only when the lambda
--     contains meaningful parallelism.
--
-- One exception is nonuniformity. A statement whose result shape varies
-- across the surrounding map-nest (e.g. 'iota x' for a mapped 'x') cannot
-- in general be traversed sequentially per thread, as arrays produced as
-- results of sequentially executed groups must be regular - handling it
-- is exactly what flattening is for. Hence irregular results still count
-- as meaningful, with one refinement in (1): if the irregular arrays are
-- used only *within* the scalar group (e.g. slices feeding a concatenation
-- whose size is ultimately uniform), the statements can be executed
-- sequentially after all - see 'mustDistribute' in 'classifyStms'. In (2)
-- irregular results make versioning worthwhile
-- ('lambdaHasMeaningfulParallelism' in Futhark.Pass.Flatten.Incremental).
--
-- The refinement above does not extend to nonuniformity arising *inside* a
-- sequential Loop or Match: an array sized by a loop-variant name implies an
-- allocation whose size cannot be computed before a kernel is launched, so a
-- statement containing one can never be part of a sequentially executed
-- scalar group, even if nothing irregular escapes it. Such statements are
-- unconditionally forced by 'mustDistribute' - see 'stmHasNonuniformInside'.
--
-- Another exception is when a trivial statement (a basic operation) produces
-- the result of a body - in that case we also treat it as meaningful, because
-- we have to manifest it. This does not extend to compound statements (e.g. a
-- sequential loop) producing a body result: those contain no parallelism worth
-- distributing on their own, and treating them as meaningful would needlessly
-- split any surrounding sequential statements into multiple kernels. This is
-- the part that would benefit from a cost model, or simply from being more
-- principled and ignoring the small inefficiencies.