-- | General definitions used for incremental flattening.
--
-- The idea behind incremental flattening is the observation that when
-- flattening a program
--
-- @
--   map f xs
-- @
--
-- we have two options: (i) transform @f@ to exploit any parallelism it may
-- contain, or (ii) turn @f@ into sequential code and only exploit the
-- parallelism in @map@.
--
-- In some cases we do not have a choice, e.g. if @f@ contains sufficiently
-- nonuniform operations that would result in nonuniform allocations. In other
-- cases the choice is obvious, such as when @f@ is completely scalar. However,
-- in the general case either will work, and it depends on the workload which of
-- the options is optimal: if the outer @map@ is big enough, it may be best to
-- efficiently sequentialise @f@ (which can then also permit various locality
-- optimisations, such as tiling). But if the outer @map@ does not have many
-- iterations, then we also need the parallelism in @f@ to fully saturate the
-- machine.
--
-- The idea behind incremental flattening is to generate both versions, and
-- select the appropriate one at run-time:
--
-- @
-- if predicate then sequentialise f...
--                   else parallelise f...
-- @
--
-- The predicate is based on comparing the amount of exploitable parallelism
-- with a threshold parameter. This threshold parameter is given a default value
-- based on run-time hardware characteristics, but usually has to be auto-tuned
-- in order to be optimal for a specific machine, program, and workload.
--
-- The multi-versioning approach is also used to generate more exotic versions,
-- such as one that parallelises @f@ at a deeper hardware level
-- (@Futhark.Pass.Flatten.Intrablock@).
module Futhark.Pass.Flatten.Incremental
  ( worthIntrablock,
    worthSequentialising,
    isVersionableMap,
    sufficientParallelism,
    isParallelFunInside,
    kernelAlternatives,
    intraBlockAlternative,
    mapAlternatives,
    scanRedAlternatives,
    propagateVersioningAttrs,
    imposeAttrs,

    -- * Transforming code
    factorScremaForParallelism,

    -- * Levels
    defaultSegLevel,
    inBlockSegLevel,
    allowVersioning,

    -- * Various queries
    bodyHasParallelism,
    lambdaHasParallelism,
    mayExploitOuter,
    onlyExploitIntra,
    mayExploitIntra,
  )
where

import Control.Monad
import Control.Monad.State
import Data.Foldable
import Data.Maybe (isJust)
import Futhark.IR.GPU
import Futhark.IR.SOACS
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Pass.Flatten.Intrablock qualified as Intrablock
import Futhark.Pass.Flatten.PreProcess
import Futhark.Tools
import Futhark.Transform.Rename
import Prelude hiding (div, quot, rem)

defaultSegLevel :: SegLevel
defaultSegLevel :: SegLevel
defaultSegLevel = SegVirt -> Maybe KernelGrid -> SegLevel
SegThread SegVirt
SegVirt Maybe KernelGrid
forall a. Maybe a
Nothing

inBlockSegLevel :: SegLevel
inBlockSegLevel :: SegLevel
inBlockSegLevel = SegVirt -> SegLevel
SegThreadInBlock SegVirt
SegNoVirt

allowVersioning :: SegLevel -> Bool
allowVersioning :: SegLevel -> Bool
allowVersioning SegThreadInBlock {} = Bool
False
allowVersioning SegLevel
_ = Bool
True

kernelAlternatives ::
  Name ->
  [Type] ->
  Body GPU ->
  [(SubExp, Body GPU)] ->
  FlattenM [VName]
kernelAlternatives :: Name
-> [Type] -> Body GPU -> [(SubExp, Body GPU)] -> FlattenM [VName]
kernelAlternatives Name
desc [Type]
_ Body GPU
default_body [] = do
  ses <- Body (Rep FlattenM) -> FlattenM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep FlattenM)
Body GPU
default_body
  forM ses $ \(SubExpRes Certs
cs SubExp
se) ->
    Certs -> FlattenM VName -> FlattenM VName
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (FlattenM VName -> FlattenM VName)
-> FlattenM VName -> FlattenM VName
forall a b. (a -> b) -> a -> b
$
      Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
desc (Exp (Rep FlattenM) -> FlattenM VName)
-> Exp (Rep FlattenM) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$
        BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$
          SubExp -> BasicOp
SubExp SubExp
se
kernelAlternatives Name
desc [Type]
result_ts Body GPU
default_body ((SubExp
cond, Body GPU
alt) : [(SubExp, Body GPU)]
alts) = do
  fallback_body <- do
    (fallback_vs, fallback_stms) <-
      FlattenM [VName] -> FlattenM ([VName], Stms (Rep FlattenM))
forall a. FlattenM a -> FlattenM (a, Stms (Rep FlattenM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms (FlattenM [VName] -> FlattenM ([VName], Stms (Rep FlattenM)))
-> FlattenM [VName] -> FlattenM ([VName], Stms (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$
        Name
-> [Type] -> Body GPU -> [(SubExp, Body GPU)] -> FlattenM [VName]
kernelAlternatives Name
desc [Type]
result_ts Body GPU
default_body [(SubExp, Body GPU)]
alts
    pure $ mkBody fallback_stms $ varsRes fallback_vs

  letTupExp desc $
    Match [cond] [Case [Just $ BoolValue True] alt] fallback_body $
      MatchDec (staticShapes result_ts) MatchEquiv

cmpSizeLe ::
  Name ->
  SizeClass ->
  [SubExp] ->
  FlattenM (SubExp, Name)
cmpSizeLe :: Name -> SizeClass -> [SubExp] -> FlattenM (SubExp, Name)
cmpSizeLe Name
desc SizeClass
size_class [SubExp]
to_what = do
  x <- (FlattenState -> Int) -> FlattenM Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets FlattenState -> Int
stateThresholdCounter
  modify $ \FlattenState
s -> FlattenState
s {stateThresholdCounter = x + 1}
  let size_key = Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> String -> Name
nameFromString (Int -> String
forall a. Show a => a -> String
show Int
x)
  to_what' <-
    letSubExp "comparatee"
      =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) to_what
  cmp_res <- letSubExp desc $ Op $ SizeOp $ CmpSizeLe size_key size_class to_what'
  pure (cmp_res, size_key)

sufficientParallelism ::
  Name ->
  [SubExp] ->
  KernelPath ->
  Maybe Int64 ->
  FlattenM (SubExp, Name)
sufficientParallelism :: Name
-> [SubExp] -> KernelPath -> Maybe Int64 -> FlattenM (SubExp, Name)
sufficientParallelism Name
desc [SubExp]
ws KernelPath
path Maybe Int64
def =
  Name -> SizeClass -> [SubExp] -> FlattenM (SubExp, Name)
cmpSizeLe Name
desc (KernelPath -> Maybe Int64 -> SizeClass
SizeThreshold KernelPath
path Maybe Int64
def) [SubExp]
ws

-- Check if the in the body there is a call to a parallel function.
-- XXX: we use this function to even reject the intra version of
-- maps that call parallel function. We should do better there.
-- One other things to note is that maybe we should create a sequential
-- version of function and replace them in these cases.
isParallelFunInside :: FunHasParallelism -> Body SOACS -> Bool
isParallelFunInside :: FunHasParallelism -> Body SOACS -> Bool
isParallelFunInside FunHasParallelism
funHasParallelism = Body SOACS -> Bool
inBody
  where
    inLambda :: GLambda SOACS t -> Bool
    inLambda :: forall t. GLambda SOACS t -> Bool
inLambda = (Stm SOACS -> Bool) -> Seq (Stm SOACS) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Exp SOACS -> Bool
callParallelFunction (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) (Seq (Stm SOACS) -> Bool)
-> (GLambda SOACS t -> Seq (Stm SOACS)) -> GLambda SOACS t -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body SOACS -> Seq (Stm SOACS)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Seq (Stm SOACS))
-> (GLambda SOACS t -> Body SOACS)
-> GLambda SOACS t
-> Seq (Stm SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GLambda SOACS t -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody
    inBody :: Body SOACS -> Bool
inBody = (Stm SOACS -> Bool) -> Seq (Stm SOACS) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Exp SOACS -> Bool
callParallelFunction (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) (Seq (Stm SOACS) -> Bool)
-> (Body SOACS -> Seq (Stm SOACS)) -> Body SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body SOACS -> Seq (Stm SOACS)
forall rep res. GBody rep res -> Stms rep
bodyStms
    callParallelFunction :: Exp SOACS -> Bool
callParallelFunction (Apply Name
fname [(SubExp, Diet)]
_ [(RetType SOACS, RetAls)]
_ Safety
_) = FunHasParallelism
funHasParallelism Name
fname
    callParallelFunction (BasicOp BasicOp
_) = Bool
False
    callParallelFunction (Match [SubExp]
_ [Case (Body SOACS)]
cases Body SOACS
def_case MatchDec (BranchType SOACS)
_) =
      Body SOACS -> Bool
inBody Body SOACS
def_case
        Bool -> Bool -> Bool
|| (Case (Body SOACS) -> Bool) -> [Case (Body SOACS)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Body SOACS -> Bool
inBody (Body SOACS -> Bool)
-> (Case (Body SOACS) -> Body SOACS) -> Case (Body SOACS) -> Bool
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
    callParallelFunction (Loop [(FParam SOACS, SubExp)]
_ LoopForm
_ Body SOACS
body) = Body SOACS -> Bool
inBody Body SOACS
body
    callParallelFunction (WithAcc [WithAccInput SOACS]
_ Lambda SOACS
lam) = Lambda SOACS -> Bool
forall t. GLambda SOACS t -> Bool
inLambda Lambda SOACS
lam
    callParallelFunction (Op (Stream SubExp
_ [VName]
_ [SubExp]
_ Lambda SOACS
lam)) = Lambda SOACS -> Bool
forall t. GLambda SOACS t -> Bool
inLambda Lambda SOACS
lam
    callParallelFunction (Op (Screma SubExp
_ [VName]
_ (ScremaForm Lambda SOACS
lam [Scan SOACS]
_ [Reduce SOACS]
_ Lambda SOACS
_))) = Lambda SOACS -> Bool
forall t. GLambda SOACS t -> Bool
inLambda Lambda SOACS
lam
    callParallelFunction (Op (Hist SubExp
_ [VName]
_ [HistOp SOACS]
ops Lambda SOACS
lam)) =
      Lambda SOACS -> Bool
forall t. GLambda SOACS t -> Bool
inLambda Lambda SOACS
lam Bool -> Bool -> Bool
|| (HistOp SOACS -> Bool) -> [HistOp SOACS] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Lambda SOACS -> Bool
forall t. GLambda SOACS t -> Bool
inLambda (Lambda SOACS -> Bool)
-> (HistOp SOACS -> Lambda SOACS) -> HistOp SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp SOACS -> Lambda SOACS
forall {rep}. HistOp rep -> Lambda rep
histLambda) [HistOp SOACS]
ops
      where
        histLambda :: HistOp rep -> Lambda rep
histLambda (Futhark.IR.SOACS.HistOp ShapeBase SubExp
_ SubExp
_ [VName]
_ [SubExp]
_ Lambda rep
op) = Lambda rep
op
    callParallelFunction (Op (FlatMap SubExp
_ [VName]
_ ExtLambda SOACS
lam)) = ExtLambda SOACS -> Bool
forall t. GLambda SOACS t -> Bool
inLambda ExtLambda SOACS
lam
    callParallelFunction (Op JVP {}) = String -> Bool
forall a. HasCallStack => String -> a
error String
"isParallelFunInside: unexpected JVP"
    callParallelFunction (Op VJP {}) = String -> Bool
forall a. HasCallStack => String -> a
error String
"isParallelFunInside: unexpected VJP"
    callParallelFunction (Op WithVJP {}) = String -> Bool
forall a. HasCallStack => String -> a
error String
"isParallelFunInside: unexpected WithVJP"

-- | Should we generate multiple versions for this map? This requires both that
-- we are at a level where versioning is possible ('allowVersioning') and that
-- the map itself produces only regular results (from an invariant width) and
-- does not call any parallel function (which force full flattening).
isVersionableMap :: FunHasParallelism -> SegLevel -> DistInputs -> DistEnv -> SubExp -> [DistResult] -> Lambda SOACS -> Bool
isVersionableMap :: FunHasParallelism
-> SegLevel
-> DistInputs
-> DistEnv
-> SubExp
-> [DistResult]
-> Lambda SOACS
-> Bool
isVersionableMap FunHasParallelism
funHasParallelism SegLevel
lvl DistInputs
inps DistEnv
_env SubExp
w [DistResult]
dist_res Lambda SOACS
map_lam =
  SegLevel -> Bool
allowVersioning SegLevel
lvl
    Bool -> Bool -> Bool
&& (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
dist_res
    Bool -> Bool -> Bool
&& Bool -> Bool
not (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
w)
    Bool -> Bool -> Bool
&& Bool -> Bool
not (FunHasParallelism -> Body SOACS -> Bool
isParallelFunInside FunHasParallelism
funHasParallelism (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
map_lam))

-- | Retrieve only those attributes that apply to flattening.
flatteningAttrs :: Attrs -> Attrs
flatteningAttrs :: Attrs -> Attrs
flatteningAttrs = [Attrs] -> Attrs
forall a. Monoid a => [a] -> a
mconcat ([Attrs] -> Attrs) -> (Attrs -> [Attrs]) -> Attrs -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Attr -> Attrs) -> Attrs -> [Attrs]
forall a. (Attr -> a) -> Attrs -> [a]
mapAttrs Attr -> Attrs
p
  where
    p :: Attr -> Attrs
p (AttrComp Name
"incremental_flattening" [Attr
x]) = Attr -> Attrs
oneAttr Attr
x
    p (AttrComp Name
"flattening" [Attr
x]) = Attr -> Attrs
oneAttr Attr
x
    p Attr
_ = Attrs
forall a. Monoid a => a
mempty

onlyExploitIntra :: Attrs -> Bool
onlyExploitIntra :: Attrs -> Bool
onlyExploitIntra Attrs
attrs =
  Attr
"only_intra" Attr -> Attrs -> Bool
`inAttrs` Attrs -> Attrs
flatteningAttrs Attrs
attrs

mayExploitOuter :: Attrs -> Bool
mayExploitOuter :: Attrs -> Bool
mayExploitOuter Attrs
attrs =
  Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Attr
"no_outer" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs' Bool -> Bool -> Bool
|| Attr
"only_inner" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs'
  where
    attrs' :: Attrs
attrs' = Attrs -> Attrs
flatteningAttrs Attrs
attrs

mayExploitIntra :: Attrs -> Bool
mayExploitIntra :: Attrs -> Bool
mayExploitIntra Attrs
attrs =
  Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Attr
"no_intra" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs' Bool -> Bool -> Bool
|| Attr
"only_inner" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs'
  where
    attrs' :: Attrs
attrs' = Attrs -> Attrs
flatteningAttrs Attrs
attrs

intraBlockAlternative ::
  Intrablock.IntrablockResult ->
  FlattenM (SubExp, Body GPU)
intraBlockAlternative :: IntrablockResult -> FlattenM (SubExp, Body GPU)
intraBlockAlternative IntrablockResult
intra = do
  Stms (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms (Stms (Rep FlattenM) -> FlattenM ())
-> Stms (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ IntrablockResult -> Stms GPU
Intrablock.intraPreludeStms IntrablockResult
intra
  max_tblock_size <-
    Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"max_tblock_size" (Exp (Rep FlattenM) -> FlattenM SubExp)
-> Exp (Rep FlattenM) -> FlattenM SubExp
forall a b. (a -> b) -> a -> b
$ Op (Rep FlattenM) -> Exp (Rep FlattenM)
forall rep. Op rep -> Exp rep
Op (Op (Rep FlattenM) -> Exp (Rep FlattenM))
-> Op (Rep FlattenM) -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$ SizeOp -> HostOp SOAC GPU
forall (op :: * -> *) rep. SizeOp -> HostOp op rep
SizeOp (SizeOp -> HostOp SOAC GPU) -> SizeOp -> HostOp SOAC GPU
forall a b. (a -> b) -> a -> b
$ SizeClass -> SizeOp
GetSizeMax SizeClass
SizeThreadBlock
  fits <-
    letSubExp "fits" $
      BasicOp $
        CmpOp
          (CmpSle Int64)
          (Intrablock.intraThreadBlockSize intra)
          max_tblock_size
  (intra_suff, _) <-
    sufficientParallelism
      "suff_intra_par"
      [Intrablock.intraAvailPar intra]
      mempty
      (Just Intrablock.intraMinInnerPar)
  intra_ok <-
    letSubExp "intra_suff_and_fits" $
      BasicOp $
        BinOp LogAnd fits intra_suff
  intra_body <-
    renameBody $
      mkBody
        (Intrablock.intraKernelStms intra)
        (varsRes $ Intrablock.intraResultNames intra)
  pure (intra_ok, intra_body)

-- | Construct the multi-versioned alternatives for a map, given actions that
-- construct the fully-flattened body and the outer-parallel-only body, and an
-- optional intrablock result. This is the shared versioning policy used both
-- for top-level maps and for maps nested inside a map-nest; the only
-- differences between the two are which bodies are supplied and how their
-- results are consumed, both of which are handled by the caller. The @ws@ are
-- the widths whose product bounds the outer parallelism (used for the threshold
-- comparison). Returns the names bound to the final results.
mapAlternatives ::
  -- | Description for the result bindings.
  Name ->
  [Type] ->
  Attrs ->
  -- | Does the map body call a parallel function? If so we must fully flatten.
  Bool ->
  -- | Is the body worth sequentialising (offering an outer-only version)?
  Bool ->
  [SubExp] ->
  -- | Construct the fully flattened body.
  FlattenM (Body GPU) ->
  -- | Construct the outer-parallelism-only body.
  FlattenM (Body GPU) ->
  Maybe Intrablock.IntrablockResult ->
  FlattenM [VName]
mapAlternatives :: Name
-> [Type]
-> Attrs
-> Bool
-> Bool
-> [SubExp]
-> FlattenM (Body GPU)
-> FlattenM (Body GPU)
-> Maybe IntrablockResult
-> FlattenM [VName]
mapAlternatives Name
desc [Type]
result_ts Attrs
attrs Bool
parallel_fun_inside Bool
worth_seq [SubExp]
ws FlattenM (Body GPU)
mkFullBody FlattenM (Body GPU)
mkOuterBody Maybe IntrablockResult
intra' =
  case Maybe IntrablockResult
intra' of
    Maybe IntrablockResult
_
      | Bool
parallel_fun_inside ->
          [(SubExp, Body GPU)] -> FlattenM [VName]
full []
      | Attr
"sequential_inner" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs -> do
          outer_body <- FlattenM (Body GPU)
mkOuterBody
          kernelAlternatives desc result_ts outer_body []
    Maybe IntrablockResult
Nothing
      | Bool -> Bool
not Bool
only_intra,
        Bool
worth_seq,
        Attrs -> Bool
mayExploitOuter Attrs
attrs ->
          [(SubExp, Body GPU)] -> FlattenM [VName]
full ([(SubExp, Body GPU)] -> FlattenM [VName])
-> ((SubExp, Body GPU) -> [(SubExp, Body GPU)])
-> (SubExp, Body GPU)
-> FlattenM [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Body GPU) -> [(SubExp, Body GPU)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((SubExp, Body GPU) -> FlattenM [VName])
-> FlattenM (SubExp, Body GPU) -> FlattenM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< FlattenM (SubExp, Body GPU)
outerAlternative
      | Bool
otherwise ->
          [(SubExp, Body GPU)] -> FlattenM [VName]
full []
    Just IntrablockResult
intra_res
      | Bool
only_intra -> do
          (_, intra_body) <- IntrablockResult -> FlattenM (SubExp, Body GPU)
intraBlockAlternative IntrablockResult
intra_res
          kernelAlternatives desc result_ts intra_body []
      | Bool
worth_seq,
        Attrs -> Bool
mayExploitOuter Attrs
attrs -> do
          outer_alt <- FlattenM (SubExp, Body GPU)
outerAlternative
          intra_alt <- intraBlockAlternative intra_res
          full [outer_alt, intra_alt]
      | Bool
otherwise ->
          [(SubExp, Body GPU)] -> FlattenM [VName]
full ([(SubExp, Body GPU)] -> FlattenM [VName])
-> ((SubExp, Body GPU) -> [(SubExp, Body GPU)])
-> (SubExp, Body GPU)
-> FlattenM [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Body GPU) -> [(SubExp, Body GPU)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((SubExp, Body GPU) -> FlattenM [VName])
-> FlattenM (SubExp, Body GPU) -> FlattenM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IntrablockResult -> FlattenM (SubExp, Body GPU)
intraBlockAlternative IntrablockResult
intra_res
  where
    only_intra :: Bool
only_intra = Attrs -> Bool
onlyExploitIntra Attrs
attrs

    full :: [(SubExp, Body GPU)] -> FlattenM [VName]
full [(SubExp, Body GPU)]
alts = do
      full_body <- FlattenM (Body GPU)
mkFullBody
      kernelAlternatives desc result_ts full_body alts

    outerAlternative :: FlattenM (SubExp, Body GPU)
outerAlternative = do
      outer_body <- FlattenM (Body GPU)
mkOuterBody
      (outer_suff, _) <- sufficientParallelism suffOuterPar ws mempty Nothing
      pure (outer_suff, outer_body)

-- | Construct the multi-versioned alternatives for a scan or reduce, given
-- actions that construct the fully-flattened body and the outer-parallel-only
-- body. Unlike 'mapAlternatives' there is no intrablock version, and the
-- outer-only version is always offered (subject to attributes).
scanRedAlternatives ::
  Name ->
  [Type] ->
  Attrs ->
  -- | Does the operator body call a parallel function? If so we must fully flatten.
  Bool ->
  -- | Does the seg level permit versioning at all (false in-block)?
  Bool ->
  [SubExp] ->
  -- | Construct the fully flattened body.
  FlattenM (Body GPU) ->
  -- | Construct the outer-parallelism-only body.
  FlattenM (Body GPU) ->
  FlattenM [VName]
scanRedAlternatives :: Name
-> [Type]
-> Attrs
-> Bool
-> Bool
-> [SubExp]
-> FlattenM (Body GPU)
-> FlattenM (Body GPU)
-> FlattenM [VName]
scanRedAlternatives Name
desc [Type]
result_ts Attrs
attrs Bool
parallel_fun_inside Bool
allow_versioning [SubExp]
ws FlattenM (Body GPU)
mkFullBody FlattenM (Body GPU)
mkOuterBody
  | Bool
parallel_fun_inside =
      FlattenM [VName]
fullAlternative
  | Attr
"sequential_inner" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs =
      FlattenM [VName]
outerAlternative
  | Attrs -> Bool
mayExploitOuter Attrs
attrs Bool -> Bool -> Bool
&& Bool
allow_versioning =
      FlattenM [VName]
fullWithOuterAlternative
  | Bool
otherwise =
      FlattenM [VName]
fullAlternative
  where
    fullAlternative :: FlattenM [VName]
fullAlternative = do
      full_body <- FlattenM (Body GPU)
mkFullBody
      kernelAlternatives desc result_ts full_body []

    outerAlternative :: FlattenM [VName]
outerAlternative = do
      outer_body <- FlattenM (Body GPU)
mkOuterBody
      kernelAlternatives desc result_ts outer_body []

    fullWithOuterAlternative :: FlattenM [VName]
fullWithOuterAlternative = do
      outer_body <- FlattenM (Body GPU)
mkOuterBody
      full_body <- mkFullBody
      (outer_suff, _) <- sufficientParallelism suffOuterPar ws mempty Nothing
      kernelAlternatives desc result_ts full_body [(outer_suff, outer_body)]

-- | The name of the threshold parameter that is used to select outer-only
-- parallelism.
suffOuterPar :: Name
suffOuterPar :: Name
suffOuterPar = Name
"suff_outer_par"

-- | Intra-group parallelism is worthwhile if the lambda contains more
-- than one instance of non-map nested parallelism, or any nested
-- parallelism inside a loop.
worthIntrablock :: Lambda SOACS -> Bool
worthIntrablock :: Lambda SOACS -> Bool
worthIntrablock Lambda SOACS
lam =
  Body SOACS -> Int
forall {rep}. (OpC rep ~ SOAC) => GBody rep SubExpRes -> Int
bodyInterest (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
  where
    bodyInterest :: GBody rep SubExpRes -> Int
bodyInterest GBody rep SubExpRes
body =
      Seq Int -> Int
forall a. Num a => Seq a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Seq Int -> Int) -> Seq Int -> Int
forall a b. (a -> b) -> a -> b
$ Stm rep -> Int
interest (Stm rep -> Int) -> Seq (Stm rep) -> Seq Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GBody rep SubExpRes -> Seq (Stm rep)
forall rep res. GBody rep res -> Stms rep
bodyStms GBody rep SubExpRes
body
    interest :: Stm rep -> Int
interest Stm rep
stm
      | Attr
"sequential" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs =
          Int
0 :: Int
      | Op (Screma SubExp
w [VName]
_ ScremaForm rep
form) <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm,
        Just Lambda rep
lam' <- ScremaForm rep -> Maybe (Lambda rep)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm rep
form =
          SubExp -> Lambda rep -> Int
forall {t}. SubExp -> GLambda rep t -> Int
mapLike SubExp
w Lambda rep
lam'
      | Loop [(FParam rep, SubExp)]
_ LoopForm
_ GBody rep SubExpRes
body <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          GBody rep SubExpRes -> Int
bodyInterest GBody rep SubExpRes
body Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
10
      | Match [SubExp]
_ [Case (GBody rep SubExpRes)]
cases GBody rep SubExpRes
defbody MatchDec (BranchType rep)
_ <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          (Int -> Int -> Int) -> Int -> [Int] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
            Int -> Int -> Int
forall a. Ord a => a -> a -> a
max
            (GBody rep SubExpRes -> Int
bodyInterest GBody rep SubExpRes
defbody)
            ((Case (GBody rep SubExpRes) -> Int)
-> [Case (GBody rep SubExpRes)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (GBody rep SubExpRes -> Int
bodyInterest (GBody rep SubExpRes -> Int)
-> (Case (GBody rep SubExpRes) -> GBody rep SubExpRes)
-> Case (GBody rep SubExpRes)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (GBody rep SubExpRes) -> GBody rep SubExpRes
forall body. Case body -> body
caseBody) [Case (GBody rep SubExpRes)]
cases)
      | Op (Screma SubExp
w [VName]
_ (ScremaForm Lambda rep
lam' [Scan rep]
_ [Reduce rep]
_ Lambda rep
_)) <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          SubExp -> Int
forall {a}. Num a => SubExp -> a
zeroIfTooSmall SubExp
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ GBody rep SubExpRes -> Int
bodyInterest (Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam')
      | Op (Stream SubExp
_ [VName]
_ [SubExp]
_ Lambda rep
lam') <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          GBody rep SubExpRes -> Int
bodyInterest (GBody rep SubExpRes -> Int) -> GBody rep SubExpRes -> Int
forall a b. (a -> b) -> a -> b
$ Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam'
      | WithAcc [WithAccInput rep]
_ Lambda rep
lam' <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          GBody rep SubExpRes -> Int
bodyInterest (GBody rep SubExpRes -> Int) -> GBody rep SubExpRes -> Int
forall a b. (a -> b) -> a -> b
$ Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam'
      | Bool
otherwise =
          Int
0
      where
        attrs :: Attrs
attrs = StmAux (ExpDec rep) -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs (StmAux (ExpDec rep) -> Attrs) -> StmAux (ExpDec rep) -> Attrs
forall a b. (a -> b) -> a -> b
$ Stm rep -> StmAux (ExpDec rep)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm rep
stm
        sequential_inner :: Bool
sequential_inner = Attr
"sequential_inner" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs

        zeroIfTooSmall :: SubExp -> a
zeroIfTooSmall (Constant (IntValue IntValue
x))
          | IntValue -> Int64
intToInt64 IntValue
x Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
32 = a
0
        zeroIfTooSmall SubExp
_ = a
1

        mapLike :: SubExp -> GLambda rep t -> Int
mapLike SubExp
w GLambda rep t
lam' =
          if Bool
sequential_inner
            then Int
0
            else Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (SubExp -> Int
forall {a}. Num a => SubExp -> a
zeroIfTooSmall SubExp
w) (GBody rep SubExpRes -> Int
bodyInterest (GLambda rep t -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda rep t
lam'))

-- | A lambda is worth sequentialising if it contains enough nested parallelism
-- of an interesting kind, or if distributing it would fragment sequential
-- control flow - that is, if it contains meaningful parallelism nested inside a
-- sequential loop or branch. Distribution must then split the loop or branch
-- into separate kernel launches (and possibly host-evaluated control flow) per
-- sequential step, so a version that instead sequentialises the nested
-- parallelism is always worth offering.
worthSequentialising :: Lambda SOACS -> Bool
worthSequentialising :: Lambda SOACS -> Bool
worthSequentialising Lambda SOACS
lam =
  Int -> Body SOACS -> Int
forall {rep} {a}.
(OpC rep ~ SOAC, Num a, Eq a) =>
a -> GBody rep SubExpRes -> Int
bodyInterest (Int
0 :: Int) (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
  where
    bodyInterest :: a -> GBody rep SubExpRes -> Int
bodyInterest a
depth GBody rep SubExpRes
body =
      Seq Int -> Int
forall a. Num a => Seq a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Seq Int -> Int) -> Seq Int -> Int
forall a b. (a -> b) -> a -> b
$ a -> Stm rep -> Int
interest a
depth (Stm rep -> Int) -> Seq (Stm rep) -> Seq Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GBody rep SubExpRes -> Seq (Stm rep)
forall rep res. GBody rep res -> Stms rep
bodyStms GBody rep SubExpRes
body
    interest :: a -> Stm rep -> Int
interest a
depth Stm rep
stm
      | Attr
"sequential" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs =
          Int
0 :: Int
      | Op (Screma SubExp
_ [VName]
_ form :: ScremaForm rep
form@(ScremaForm Lambda rep
lam' [Scan rep]
_ [Reduce rep]
_ Lambda rep
_)) <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm,
        Maybe (Lambda rep) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Lambda rep) -> Bool) -> Maybe (Lambda rep) -> Bool
forall a b. (a -> b) -> a -> b
$ ScremaForm rep -> Maybe (Lambda rep)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm rep
form =
          if Bool
sequential_inner
            then Int
0
            else a -> GBody rep SubExpRes -> Int
bodyInterest (a
depth a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam')
      | Loop [(FParam rep, SubExp)]
_ LoopForm
_ GBody rep SubExpRes
body <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          a -> GBody rep SubExpRes -> Int
bodyInterest (a
depth a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) GBody rep SubExpRes
body Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
10
      | Match [SubExp]
_ [Case (GBody rep SubExpRes)]
cases GBody rep SubExpRes
defbody MatchDec (BranchType rep)
_ <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
*) (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$
            [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$
              (GBody rep SubExpRes -> Int) -> [GBody rep SubExpRes] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (a -> GBody rep SubExpRes -> Int
bodyInterest (a
depth a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)) ([GBody rep SubExpRes] -> [Int]) -> [GBody rep SubExpRes] -> [Int]
forall a b. (a -> b) -> a -> b
$
                GBody rep SubExpRes
defbody GBody rep SubExpRes
-> [GBody rep SubExpRes] -> [GBody rep SubExpRes]
forall a. a -> [a] -> [a]
: (Case (GBody rep SubExpRes) -> GBody rep SubExpRes)
-> [Case (GBody rep SubExpRes)] -> [GBody rep SubExpRes]
forall a b. (a -> b) -> [a] -> [b]
map Case (GBody rep SubExpRes) -> GBody rep SubExpRes
forall body. Case body -> body
caseBody [Case (GBody rep SubExpRes)]
cases
      | WithAcc [WithAccInput rep]
_ Lambda rep
withacc_lam <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          a -> GBody rep SubExpRes -> Int
bodyInterest (a
depth a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
withacc_lam)
      | Op (Screma SubExp
_ [VName]
_ form :: ScremaForm rep
form@(ScremaForm Lambda rep
lam' [Scan rep]
_ [Reduce rep]
_ Lambda rep
_)) <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          Int
1
            Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> GBody rep SubExpRes -> Int
bodyInterest (a
depth a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam')
            Int -> Int -> Int
forall a. Num a => a -> a -> a
+
            -- Give this a bigger score if it's a redomap just inside
            -- the the outer lambda, as these are often tileable and
            -- thus benefit more from sequentialisation.
            case (ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm rep
form, a
depth) of
              (Just ([Reduce rep], Lambda rep)
_, a
0) -> Int
1
              (Maybe ([Reduce rep], Lambda rep), a)
_ -> Int
0
      | Op (Stream SubExp
_ [VName]
_ [SubExp]
_ Lambda rep
lam') <- Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm =
          a -> GBody rep SubExpRes -> Int
bodyInterest (a
depth a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
lam')
      | Bool
otherwise =
          Int
0
      where
        attrs :: Attrs
attrs = StmAux (ExpDec rep) -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs (StmAux (ExpDec rep) -> Attrs) -> StmAux (ExpDec rep) -> Attrs
forall a b. (a -> b) -> a -> b
$ Stm rep -> StmAux (ExpDec rep)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm rep
stm
        sequential_inner :: Bool
sequential_inner = Attr
"sequential_inner" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs

bodyHasParallelism :: FunHasParallelism -> Body SOACS -> Bool
bodyHasParallelism :: FunHasParallelism -> Body SOACS -> Bool
bodyHasParallelism FunHasParallelism
funHasParallelism =
  (Stm SOACS -> Bool) -> Seq (Stm SOACS) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (FunHasParallelism -> Stm SOACS -> Bool
isParallelStm FunHasParallelism
funHasParallelism) (Seq (Stm SOACS) -> Bool)
-> (Body SOACS -> Seq (Stm SOACS)) -> Body SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body SOACS -> Seq (Stm SOACS)
forall rep res. GBody rep res -> Stms rep
bodyStms

lambdaHasParallelism :: FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasParallelism :: FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasParallelism FunHasParallelism
funHasParallelism =
  FunHasParallelism -> Body SOACS -> Bool
bodyHasParallelism FunHasParallelism
funHasParallelism (Body SOACS -> Bool)
-> (Lambda SOACS -> Body SOACS) -> Lambda SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody

-- | Like 'lambdaHasParallelism', but only counts meaningful
-- parallelism: a SOAC, a call to a parallel function, or a statement
-- with an irregular result, which requires flattening to exploit.
-- Basic operations such as 'Replicate' of invariant size do not
-- provide enough parallelism on their own to make multi-versioning
-- worthwhile.  See Note [Meaningful Parallelism] in
-- Futhark.Pass.Flatten.Distribute.
lambdaHasMeaningfulParallelism :: FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasMeaningfulParallelism :: FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasMeaningfulParallelism FunHasParallelism
funHasParallelism Lambda SOACS
lam =
  (Stm SOACS -> Bool) -> Seq (Stm SOACS) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Stm SOACS -> Bool
interesting (Seq (Stm SOACS) -> Bool) -> Seq (Stm SOACS) -> Bool
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Seq (Stm SOACS)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Seq (Stm SOACS)) -> Body SOACS -> Seq (Stm SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
  where
    free_in_lam :: Names
free_in_lam = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
lam
    invariantDim :: SubExp -> Bool
invariantDim (Var VName
v) = VName
v VName -> Names -> Bool
`nameIn` Names
free_in_lam
    invariantDim Constant {} = Bool
True
    irregularResult :: Stm SOACS -> Bool
irregularResult =
      Bool -> Bool
not (Bool -> Bool) -> (Stm SOACS -> Bool) -> Stm SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all SubExp -> Bool
invariantDim ([SubExp] -> Bool) -> (Type -> [SubExp]) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims) ([Type] -> Bool) -> (Stm SOACS -> [Type]) -> Stm SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec SOACS) -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes (Pat (LetDec SOACS) -> [Type])
-> (Stm SOACS -> Pat (LetDec SOACS)) -> Stm SOACS -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm SOACS -> Pat (LetDec SOACS)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat
    interesting :: Stm SOACS -> Bool
interesting Stm SOACS
stm =
      FunHasParallelism -> Stm SOACS -> Bool
stmHasMeaningfulParallelism FunHasParallelism
funHasParallelism Stm SOACS
stm Bool -> Bool -> Bool
|| Stm SOACS -> Bool
irregularResult Stm SOACS
stm

-- | Produce a body suitable for full flattening from a Screma, or
-- 'Nothing' if none of its lambdas contain meaningful parallelism, in
-- which case multi-versioning is not worthwhile.  See Note
-- [Meaningful Parallelism] in Futhark.Pass.Flatten.Distribute.
factorScremaForParallelism ::
  (MonadBuilder m) =>
  FunHasParallelism ->
  Scope SOACS ->
  StmAux () ->
  Pat Type ->
  SubExp ->
  [VName] ->
  ScremaForm SOACS ->
  m (Maybe (Body SOACS))
factorScremaForParallelism :: forall (m :: * -> *).
MonadBuilder m =>
FunHasParallelism
-> Scope SOACS
-> StmAux ()
-> Pat Type
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> m (Maybe (Body SOACS))
factorScremaForParallelism FunHasParallelism
funHasParallelism Scope SOACS
scope StmAux ()
aux Pat Type
pat SubExp
w [VName]
arrs ScremaForm SOACS
form
  | Just ([Reduce SOACS]
reds, Lambda SOACS
map_lam) <- ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm SOACS
form,
    FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasMeaningfulParallelism FunHasParallelism
funHasParallelism Lambda SOACS
map_lam = do
      map_lam' <- Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
preprocessLambda Scope SOACS
scope Lambda SOACS
map_lam
      (map_stm, red_stm) <-
        redomapToMapAndReduce
          pat
          (w, reds, map_lam', arrs)
      Just <$> mkFactoredBody (stmsFromList [map_stm, red_stm])
  | Just (Lambda SOACS
post_lam, [Scan SOACS]
scans, Lambda SOACS
map_lam) <- ScremaForm SOACS
-> Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS)
forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC ScremaForm SOACS
form,
    FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasMeaningfulParallelism FunHasParallelism
funHasParallelism Lambda SOACS
map_lam,
    FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasMeaningfulParallelism FunHasParallelism
funHasParallelism Lambda SOACS
post_lam = do
      map_lam' <- Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
preprocessLambda Scope SOACS
scope Lambda SOACS
map_lam
      post_lam' <- preprocessLambda scope post_lam
      (map_stm, scan_stm, post_stm) <-
        maposcanomapToMapScanAndMap
          pat
          (w, post_lam', scans, map_lam', arrs)
      Just <$> mkFactoredBody (stmsFromList [map_stm, scan_stm, post_stm])
  | Just (Lambda SOACS
post_lam, [Scan SOACS]
scans, Lambda SOACS
map_lam) <- ScremaForm SOACS
-> Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS)
forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC ScremaForm SOACS
form,
    FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasMeaningfulParallelism FunHasParallelism
funHasParallelism Lambda SOACS
map_lam = do
      map_lam' <- Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
preprocessLambda Scope SOACS
scope Lambda SOACS
map_lam
      post_lam' <- preprocessLambda scope post_lam
      (map_stm, scanomap_stm) <-
        maposcanomapToMaposcanAndMap
          pat
          (w, post_lam', scans, map_lam', arrs)
      Just <$> mkFactoredBody (stmsFromList [map_stm, scanomap_stm])
  | Just (Lambda SOACS
post_lam, [Scan SOACS]
scans, Lambda SOACS
map_lam) <- ScremaForm SOACS
-> Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS)
forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC ScremaForm SOACS
form,
    FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasMeaningfulParallelism FunHasParallelism
funHasParallelism Lambda SOACS
post_lam = do
      map_lam' <- Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
preprocessLambda Scope SOACS
scope Lambda SOACS
map_lam
      post_lam' <- preprocessLambda scope post_lam
      (map_stm, scan_stm, post_stm) <-
        maposcanomapToMapScanAndMap
          pat
          (w, post_lam', scans, map_lam', arrs)
      Just <$> mkFactoredBody (stmsFromList [map_stm, scan_stm, post_stm])
  | Bool
otherwise =
      Maybe (Body SOACS) -> m (Maybe (Body SOACS))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Body SOACS)
forall a. Maybe a
Nothing
  where
    mkFactoredBody :: Seq (Stm SOACS) -> m (Body SOACS)
mkFactoredBody Seq (Stm SOACS)
stms = do
      stms' <-
        (Stm SOACS -> Stm SOACS) -> Seq (Stm SOACS) -> Seq (Stm SOACS)
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Attrs -> Stm SOACS -> Stm SOACS
propagateAttrs (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
aux) (Stm SOACS -> Stm SOACS)
-> (Stm SOACS -> Stm SOACS) -> Stm SOACS -> Stm SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Certs -> Stm SOACS -> Stm SOACS
forall rep. Certs -> Stm rep -> Stm rep
certify (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
aux))
          (Seq (Stm SOACS) -> Seq (Stm SOACS))
-> m (Seq (Stm SOACS)) -> m (Seq (Stm SOACS))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Scope SOACS -> Seq (Stm SOACS) -> m (Seq (Stm SOACS))
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Seq (Stm SOACS) -> m (Seq (Stm SOACS))
preprocessStms Scope SOACS
scope Seq (Stm SOACS)
stms
      pure $ mkBody stms' $ varsRes $ patNames pat

-- | Add the flattening attributes of the enclosing context to a statement. A
-- statement that carries flattening attributes of its own is left alone, as
-- those are more specific.
propagateAttrs :: Attrs -> Stm SOACS -> Stm SOACS
propagateAttrs :: Attrs -> Stm SOACS -> Stm SOACS
propagateAttrs Attrs
attrs Stm SOACS
stm
  | Attrs
attrs' Attrs -> Attrs -> Bool
forall a. Eq a => a -> a -> Bool
== Attrs
forall a. Monoid a => a
mempty = Stm SOACS
stm
  | Attrs -> Attrs
flatteningAttrs (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs (Stm SOACS -> StmAux (ExpDec SOACS)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm SOACS
stm)) Attrs -> Attrs -> Bool
forall a. Eq a => a -> a -> Bool
== Attrs
forall a. Monoid a => a
mempty =
      Stm SOACS
stm {stmAux = (stmAux stm) {stmAuxAttrs = attrs' <> stmAuxAttrs (stmAux stm)}}
  | Bool
otherwise = Stm SOACS
stm
  where
    -- 'flatteningAttrs' strips the enclosing 'flattening', which has to be put
    -- back for the attributes to be recognised on the statement.
    attrs' :: Attrs
attrs' =
      [Attrs] -> Attrs
forall a. Monoid a => [a] -> a
mconcat ([Attrs] -> Attrs) -> (Attrs -> [Attrs]) -> Attrs -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Attr -> Attrs) -> Attrs -> [Attrs]
forall a. (Attr -> a) -> Attrs -> [a]
mapAttrs (Attr -> Attrs
oneAttr (Attr -> Attrs) -> (Attr -> Attr) -> Attr -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Attr] -> Attr
AttrComp Name
"flattening" ([Attr] -> Attr) -> (Attr -> [Attr]) -> Attr -> Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Attr -> [Attr]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure) (Attrs -> Attrs) -> Attrs -> Attrs
forall a b. (a -> b) -> a -> b
$
        Attrs -> Attrs
flatteningAttrs Attrs
attrs

-- | Impose outside flattening attributes on a statement. Only SOACs are
-- affected, and only those that carry no flattening attributes of their own, as
-- those are more specific.
imposeAttrs :: Attrs -> Stm SOACS -> Stm SOACS
imposeAttrs :: Attrs -> Stm SOACS -> Stm SOACS
imposeAttrs Attrs
attrs Stm SOACS
stm
  | Op {} <- Stm SOACS -> Exp SOACS
forall rep. Stm rep -> Exp rep
stmExp Stm SOACS
stm = Attrs -> Stm SOACS -> Stm SOACS
propagateAttrs Attrs
attrs Stm SOACS
stm
  | Bool
otherwise = Stm SOACS
stm

-- | Propagate incremental flattening attributes to the statements of
-- a map lambda body. Statements that carry their own incremental
-- flattening attributes are left alone.
propagateVersioningAttrs :: Attrs -> Lambda SOACS -> Lambda SOACS
propagateVersioningAttrs :: Attrs -> Lambda SOACS -> Lambda SOACS
propagateVersioningAttrs Attrs
attrs Lambda SOACS
lam
  | Attrs -> Attrs
flatteningAttrs Attrs
attrs Attrs -> Attrs -> Bool
forall a. Eq a => a -> a -> Bool
== Attrs
forall a. Monoid a => a
mempty = Lambda SOACS
lam
  | Bool
otherwise =
      Lambda SOACS
lam
        { lambdaBody =
            (lambdaBody lam)
              { bodyStms = fmap (propagateAttrs attrs) (bodyStms (lambdaBody lam))
              }
        }