{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-overlapping-patterns -Wno-incomplete-patterns -Wno-incomplete-uni-patterns -Wno-incomplete-record-updates #-}

module Futhark.Pass.ExtractKernels.DistributeNests
  ( MapLoop (..),
    mapLoopStm,
    bodyContainsParallelism,
    lambdaContainsParallelism,
    determineReduceOp,
    histKernel,
    DistEnv (..),
    DistAcc (..),
    runDistNestT,
    DistNestT,
    liftInner,
    distributeMap,
    distribute,
    distributeSingleStm,
    distributeMapBodyStms,
    addStmsToAcc,
    addStmToAcc,
    permutationAndMissing,
    addPostStms,
    postStm,
    inNesting,
  )
where

import Control.Arrow (first)
import Control.Monad.Identity
import Control.Monad.RWS.Strict
import Control.Monad.Reader
import Control.Monad.Trans.Maybe
import Control.Monad.Writer.Strict
import Data.Function ((&))
import Data.List (find, partition, tails)
import qualified Data.Map as M
import Data.Maybe
import Futhark.IR
import Futhark.IR.SOACS (SOACS)
import qualified Futhark.IR.SOACS as SOACS
import Futhark.IR.SOACS.SOAC hiding (HistOp, histDest)
import Futhark.IR.SOACS.Simplify (simpleSOACS, simplifyStms)
import Futhark.IR.SegOp
import Futhark.MonadFreshNames
import Futhark.Pass.ExtractKernels.BlockedKernel
import Futhark.Pass.ExtractKernels.Distribution
import Futhark.Pass.ExtractKernels.ISRWIM
import Futhark.Pass.ExtractKernels.Interchange
import Futhark.Tools
import Futhark.Transform.CopyPropagate
import qualified Futhark.Transform.FirstOrderTransform as FOT
import Futhark.Transform.Rename
import Futhark.Util
import Futhark.Util.Log

scopeForSOACs :: SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs :: forall rep. SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs = Scope rep -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope

data MapLoop = MapLoop SOACS.Pattern (StmAux ()) SubExp SOACS.Lambda [VName]

mapLoopStm :: MapLoop -> Stm SOACS
mapLoopStm :: MapLoop -> Stm SOACS
mapLoopStm (MapLoop Pattern
pat StmAux ()
aux SubExp
w Lambda
lam [VName]
arrs) =
  Pattern -> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pattern
pat StmAux ()
StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ Op SOACS -> Exp SOACS
forall rep. Op rep -> ExpT rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs (ScremaForm SOACS -> SOAC SOACS) -> ScremaForm SOACS -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ Lambda -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda
lam

data DistEnv rep m = DistEnv
  { forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest :: Nestings,
    forall rep (m :: * -> *). DistEnv rep m -> Scope rep
distScope :: Scope rep,
    forall rep (m :: * -> *).
DistEnv rep m -> Stms SOACS -> DistNestT rep m (Stms rep)
distOnTopLevelStms :: Stms SOACS -> DistNestT rep m (Stms rep),
    forall rep (m :: * -> *).
DistEnv rep m
-> MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
distOnInnerMap ::
      MapLoop ->
      DistAcc rep ->
      DistNestT rep m (DistAcc rep),
    forall rep (m :: * -> *).
DistEnv rep m -> Stm SOACS -> Binder rep (Stms rep)
distOnSOACSStms :: Stm SOACS -> Binder rep (Stms rep),
    forall rep (m :: * -> *).
DistEnv rep m -> Lambda -> Binder rep (Lambda rep)
distOnSOACSLambda :: Lambda SOACS -> Binder rep (Lambda rep),
    forall rep (m :: * -> *). DistEnv rep m -> MkSegLevel rep m
distSegLevel :: MkSegLevel rep m
  }

data DistAcc rep = DistAcc
  { forall rep. DistAcc rep -> Targets
distTargets :: Targets,
    forall rep. DistAcc rep -> Stms rep
distStms :: Stms rep
  }

data DistRes rep = DistRes
  { forall rep. DistRes rep -> PostStms rep
accPostStms :: PostStms rep,
    forall rep. DistRes rep -> Log
accLog :: Log
  }

instance Semigroup (DistRes rep) where
  DistRes PostStms rep
ks1 Log
log1 <> :: DistRes rep -> DistRes rep -> DistRes rep
<> DistRes PostStms rep
ks2 Log
log2 =
    PostStms rep -> Log -> DistRes rep
forall rep. PostStms rep -> Log -> DistRes rep
DistRes (PostStms rep
ks1 PostStms rep -> PostStms rep -> PostStms rep
forall a. Semigroup a => a -> a -> a
<> PostStms rep
ks2) (Log
log1 Log -> Log -> Log
forall a. Semigroup a => a -> a -> a
<> Log
log2)

instance Monoid (DistRes rep) where
  mempty :: DistRes rep
mempty = PostStms rep -> Log -> DistRes rep
forall rep. PostStms rep -> Log -> DistRes rep
DistRes PostStms rep
forall a. Monoid a => a
mempty Log
forall a. Monoid a => a
mempty

newtype PostStms rep = PostStms {forall rep. PostStms rep -> Stms rep
unPostStms :: Stms rep}

instance Semigroup (PostStms rep) where
  PostStms Stms rep
xs <> :: PostStms rep -> PostStms rep -> PostStms rep
<> PostStms Stms rep
ys = Stms rep -> PostStms rep
forall rep. Stms rep -> PostStms rep
PostStms (Stms rep -> PostStms rep) -> Stms rep -> PostStms rep
forall a b. (a -> b) -> a -> b
$ Stms rep
ys Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Stms rep
xs

instance Monoid (PostStms rep) where
  mempty :: PostStms rep
mempty = Stms rep -> PostStms rep
forall rep. Stms rep -> PostStms rep
PostStms Stms rep
forall a. Monoid a => a
mempty

typeEnvFromDistAcc :: DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc :: forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc = PatternT Type -> Scope rep
forall rep dec. (LetDec rep ~ dec) => PatternT dec -> Scope rep
scopeOfPattern (PatternT Type -> Scope rep)
-> (DistAcc rep -> PatternT Type) -> DistAcc rep -> Scope rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PatternT Type, [SubExp]) -> PatternT Type
forall a b. (a, b) -> a
fst ((PatternT Type, [SubExp]) -> PatternT Type)
-> (DistAcc rep -> (PatternT Type, [SubExp]))
-> DistAcc rep
-> PatternT Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Targets -> (PatternT Type, [SubExp])
outerTarget (Targets -> (PatternT Type, [SubExp]))
-> (DistAcc rep -> Targets)
-> DistAcc rep
-> (PatternT Type, [SubExp])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistAcc rep -> Targets
forall rep. DistAcc rep -> Targets
distTargets

addStmsToAcc :: Stms rep -> DistAcc rep -> DistAcc rep
addStmsToAcc :: forall rep. Stms rep -> DistAcc rep -> DistAcc rep
addStmsToAcc Stms rep
stms DistAcc rep
acc =
  DistAcc rep
acc {distStms :: Stms rep
distStms = Stms rep
stms Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> DistAcc rep -> Stms rep
forall rep. DistAcc rep -> Stms rep
distStms DistAcc rep
acc}

addStmToAcc ::
  (MonadFreshNames m, DistRep rep) =>
  Stm SOACS ->
  DistAcc rep ->
  DistNestT rep m (DistAcc rep)
addStmToAcc :: forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc = do
  Stm SOACS -> Binder rep (Stms rep)
onSoacs <- (DistEnv rep m -> Stm SOACS -> Binder rep (Stms rep))
-> DistNestT rep m (Stm SOACS -> Binder rep (Stms rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Stm SOACS -> Binder rep (Stms rep)
forall rep (m :: * -> *).
DistEnv rep m -> Stm SOACS -> Binder rep (Stms rep)
distOnSOACSStms
  (Stms rep
stm', Stms rep
_) <- Binder rep (Stms rep) -> DistNestT rep m (Stms rep, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Binder rep (Stms rep) -> DistNestT rep m (Stms rep, Stms rep))
-> Binder rep (Stms rep) -> DistNestT rep m (Stms rep, Stms rep)
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Binder rep (Stms rep)
onSoacs Stm SOACS
stm
  DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc {distStms :: Stms rep
distStms = Stms rep
stm' Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> DistAcc rep -> Stms rep
forall rep. DistAcc rep -> Stms rep
distStms DistAcc rep
acc}

soacsLambda ::
  (MonadFreshNames m, DistRep rep) =>
  Lambda SOACS ->
  DistNestT rep m (Lambda rep)
soacsLambda :: forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Lambda -> DistNestT rep m (Lambda rep)
soacsLambda Lambda
lam = do
  Lambda -> Binder rep (Lambda rep)
onLambda <- (DistEnv rep m -> Lambda -> Binder rep (Lambda rep))
-> DistNestT rep m (Lambda -> Binder rep (Lambda rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Lambda -> Binder rep (Lambda rep)
forall rep (m :: * -> *).
DistEnv rep m -> Lambda -> Binder rep (Lambda rep)
distOnSOACSLambda
  (Lambda rep, Stms rep) -> Lambda rep
forall a b. (a, b) -> a
fst ((Lambda rep, Stms rep) -> Lambda rep)
-> DistNestT rep m (Lambda rep, Stms rep)
-> DistNestT rep m (Lambda rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Binder rep (Lambda rep) -> DistNestT rep m (Lambda rep, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Lambda -> Binder rep (Lambda rep)
onLambda Lambda
lam)

newtype DistNestT rep m a
  = DistNestT (ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a)
  deriving
    ( (forall a b. (a -> b) -> DistNestT rep m a -> DistNestT rep m b)
-> (forall a b. a -> DistNestT rep m b -> DistNestT rep m a)
-> Functor (DistNestT rep m)
forall a b. a -> DistNestT rep m b -> DistNestT rep m a
forall a b. (a -> b) -> DistNestT rep m a -> DistNestT rep m b
forall rep (m :: * -> *) a b.
Functor m =>
a -> DistNestT rep m b -> DistNestT rep m a
forall rep (m :: * -> *) a b.
Functor m =>
(a -> b) -> DistNestT rep m a -> DistNestT rep m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> DistNestT rep m b -> DistNestT rep m a
$c<$ :: forall rep (m :: * -> *) a b.
Functor m =>
a -> DistNestT rep m b -> DistNestT rep m a
fmap :: forall a b. (a -> b) -> DistNestT rep m a -> DistNestT rep m b
$cfmap :: forall rep (m :: * -> *) a b.
Functor m =>
(a -> b) -> DistNestT rep m a -> DistNestT rep m b
Functor,
      Functor (DistNestT rep m)
Functor (DistNestT rep m)
-> (forall a. a -> DistNestT rep m a)
-> (forall a b.
    DistNestT rep m (a -> b) -> DistNestT rep m a -> DistNestT rep m b)
-> (forall a b c.
    (a -> b -> c)
    -> DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m c)
-> (forall a b.
    DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b)
-> (forall a b.
    DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m a)
-> Applicative (DistNestT rep m)
forall a. a -> DistNestT rep m a
forall a b.
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m a
forall a b.
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
forall a b.
DistNestT rep m (a -> b) -> DistNestT rep m a -> DistNestT rep m b
forall a b c.
(a -> b -> c)
-> DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m c
forall {rep} {m :: * -> *}.
Applicative m =>
Functor (DistNestT rep m)
forall rep (m :: * -> *) a. Applicative m => a -> DistNestT rep m a
forall rep (m :: * -> *) a b.
Applicative m =>
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m a
forall rep (m :: * -> *) a b.
Applicative m =>
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
forall rep (m :: * -> *) a b.
Applicative m =>
DistNestT rep m (a -> b) -> DistNestT rep m a -> DistNestT rep m b
forall rep (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c)
-> DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b.
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m a
$c<* :: forall rep (m :: * -> *) a b.
Applicative m =>
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m a
*> :: forall a b.
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
$c*> :: forall rep (m :: * -> *) a b.
Applicative m =>
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
liftA2 :: forall a b c.
(a -> b -> c)
-> DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m c
$cliftA2 :: forall rep (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c)
-> DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m c
<*> :: forall a b.
DistNestT rep m (a -> b) -> DistNestT rep m a -> DistNestT rep m b
$c<*> :: forall rep (m :: * -> *) a b.
Applicative m =>
DistNestT rep m (a -> b) -> DistNestT rep m a -> DistNestT rep m b
pure :: forall a. a -> DistNestT rep m a
$cpure :: forall rep (m :: * -> *) a. Applicative m => a -> DistNestT rep m a
Applicative,
      Applicative (DistNestT rep m)
Applicative (DistNestT rep m)
-> (forall a b.
    DistNestT rep m a -> (a -> DistNestT rep m b) -> DistNestT rep m b)
-> (forall a b.
    DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b)
-> (forall a. a -> DistNestT rep m a)
-> Monad (DistNestT rep m)
forall a. a -> DistNestT rep m a
forall a b.
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
forall a b.
DistNestT rep m a -> (a -> DistNestT rep m b) -> DistNestT rep m b
forall {rep} {m :: * -> *}.
Monad m =>
Applicative (DistNestT rep m)
forall rep (m :: * -> *) a. Monad m => a -> DistNestT rep m a
forall rep (m :: * -> *) a b.
Monad m =>
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
forall rep (m :: * -> *) a b.
Monad m =>
DistNestT rep m a -> (a -> DistNestT rep m b) -> DistNestT rep m b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> DistNestT rep m a
$creturn :: forall rep (m :: * -> *) a. Monad m => a -> DistNestT rep m a
>> :: forall a b.
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
$c>> :: forall rep (m :: * -> *) a b.
Monad m =>
DistNestT rep m a -> DistNestT rep m b -> DistNestT rep m b
>>= :: forall a b.
DistNestT rep m a -> (a -> DistNestT rep m b) -> DistNestT rep m b
$c>>= :: forall rep (m :: * -> *) a b.
Monad m =>
DistNestT rep m a -> (a -> DistNestT rep m b) -> DistNestT rep m b
Monad,
      MonadReader (DistEnv rep m),
      MonadWriter (DistRes rep)
    )

liftInner :: (LocalScope rep m, DistRep rep) => m a -> DistNestT rep m a
liftInner :: forall rep (m :: * -> *) a.
(LocalScope rep m, DistRep rep) =>
m a -> DistNestT rep m a
liftInner m a
m = do
  Scope rep
outer_scope <- DistNestT rep m (Scope rep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
-> DistNestT rep m a
forall rep (m :: * -> *) a.
ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
-> DistNestT rep m a
DistNestT (ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
 -> DistNestT rep m a)
-> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
-> DistNestT rep m a
forall a b. (a -> b) -> a -> b
$
    WriterT (DistRes rep) m a
-> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (WriterT (DistRes rep) m a
 -> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a)
-> WriterT (DistRes rep) m a
-> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
forall a b. (a -> b) -> a -> b
$
      m a -> WriterT (DistRes rep) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> WriterT (DistRes rep) m a)
-> m a -> WriterT (DistRes rep) m a
forall a b. (a -> b) -> a -> b
$ do
        Scope rep
inner_scope <- m (Scope rep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
        Scope rep -> m a -> m a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Scope rep
outer_scope Scope rep -> Scope rep -> Scope rep
forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` Scope rep
inner_scope) m a
m

instance MonadFreshNames m => MonadFreshNames (DistNestT rep m) where
  getNameSource :: DistNestT rep m VNameSource
getNameSource = ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) VNameSource
-> DistNestT rep m VNameSource
forall rep (m :: * -> *) a.
ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
-> DistNestT rep m a
DistNestT (ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) VNameSource
 -> DistNestT rep m VNameSource)
-> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) VNameSource
-> DistNestT rep m VNameSource
forall a b. (a -> b) -> a -> b
$ WriterT (DistRes rep) m VNameSource
-> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) VNameSource
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift WriterT (DistRes rep) m VNameSource
forall (m :: * -> *). MonadFreshNames m => m VNameSource
getNameSource
  putNameSource :: VNameSource -> DistNestT rep m ()
putNameSource = ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) ()
-> DistNestT rep m ()
forall rep (m :: * -> *) a.
ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a
-> DistNestT rep m a
DistNestT (ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) ()
 -> DistNestT rep m ())
-> (VNameSource
    -> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) ())
-> VNameSource
-> DistNestT rep m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WriterT (DistRes rep) m ()
-> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (WriterT (DistRes rep) m ()
 -> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) ())
-> (VNameSource -> WriterT (DistRes rep) m ())
-> VNameSource
-> ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VNameSource -> WriterT (DistRes rep) m ()
forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
putNameSource

instance (Monad m, ASTRep rep) => HasScope rep (DistNestT rep m) where
  askScope :: DistNestT rep m (Scope rep)
askScope = (DistEnv rep m -> Scope rep) -> DistNestT rep m (Scope rep)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Scope rep
forall rep (m :: * -> *). DistEnv rep m -> Scope rep
distScope

instance (Monad m, ASTRep rep) => LocalScope rep (DistNestT rep m) where
  localScope :: forall a. Scope rep -> DistNestT rep m a -> DistNestT rep m a
localScope Scope rep
types = (DistEnv rep m -> DistEnv rep m)
-> DistNestT rep m a -> DistNestT rep m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((DistEnv rep m -> DistEnv rep m)
 -> DistNestT rep m a -> DistNestT rep m a)
-> (DistEnv rep m -> DistEnv rep m)
-> DistNestT rep m a
-> DistNestT rep m a
forall a b. (a -> b) -> a -> b
$ \DistEnv rep m
env ->
    DistEnv rep m
env {distScope :: Scope rep
distScope = Scope rep
types Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> DistEnv rep m -> Scope rep
forall rep (m :: * -> *). DistEnv rep m -> Scope rep
distScope DistEnv rep m
env}

instance Monad m => MonadLogger (DistNestT rep m) where
  addLog :: Log -> DistNestT rep m ()
addLog Log
msgs = DistRes rep -> DistNestT rep m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell DistRes rep
forall a. Monoid a => a
mempty {accLog :: Log
accLog = Log
msgs}

runDistNestT ::
  (MonadLogger m, DistRep rep) =>
  DistEnv rep m ->
  DistNestT rep m (DistAcc rep) ->
  m (Stms rep)
runDistNestT :: forall (m :: * -> *) rep.
(MonadLogger m, DistRep rep) =>
DistEnv rep m -> DistNestT rep m (DistAcc rep) -> m (Stms rep)
runDistNestT DistEnv rep m
env (DistNestT ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) (DistAcc rep)
m) = do
  (DistAcc rep
acc, DistRes rep
res) <- WriterT (DistRes rep) m (DistAcc rep)
-> m (DistAcc rep, DistRes rep)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT (DistRes rep) m (DistAcc rep)
 -> m (DistAcc rep, DistRes rep))
-> WriterT (DistRes rep) m (DistAcc rep)
-> m (DistAcc rep, DistRes rep)
forall a b. (a -> b) -> a -> b
$ ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) (DistAcc rep)
-> DistEnv rep m -> WriterT (DistRes rep) m (DistAcc rep)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) (DistAcc rep)
m DistEnv rep m
env
  Log -> m ()
forall (m :: * -> *). MonadLogger m => Log -> m ()
addLog (Log -> m ()) -> Log -> m ()
forall a b. (a -> b) -> a -> b
$ DistRes rep -> Log
forall rep. DistRes rep -> Log
accLog DistRes rep
res
  -- There may be a few final targets remaining - these correspond to
  -- arrays that are identity mapped, and must have statements
  -- inserted here.
  Stms rep -> m (Stms rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms rep -> m (Stms rep)) -> Stms rep -> m (Stms rep)
forall a b. (a -> b) -> a -> b
$
    PostStms rep -> Stms rep
forall rep. PostStms rep -> Stms rep
unPostStms (DistRes rep -> PostStms rep
forall rep. DistRes rep -> PostStms rep
accPostStms DistRes rep
res)
      Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> (PatternT Type, [SubExp]) -> Stms rep
identityStms (Targets -> (PatternT Type, [SubExp])
outerTarget (Targets -> (PatternT Type, [SubExp]))
-> Targets -> (PatternT Type, [SubExp])
forall a b. (a -> b) -> a -> b
$ DistAcc rep -> Targets
forall rep. DistAcc rep -> Targets
distTargets DistAcc rep
acc)
  where
    outermost :: LoopNesting
outermost = Nesting -> LoopNesting
nestingLoop (Nesting -> LoopNesting) -> Nesting -> LoopNesting
forall a b. (a -> b) -> a -> b
$
      case DistEnv rep m -> Nestings
forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest DistEnv rep m
env of
        (Nesting
nest, []) -> Nesting
nest
        (Nesting
_, Nesting
nest : [Nesting]
_) -> Nesting
nest
    params_to_arrs :: [(VName, VName)]
params_to_arrs =
      ((Param Type, VName) -> (VName, VName))
-> [(Param Type, VName)] -> [(VName, VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((Param Type -> VName) -> (Param Type, VName) -> (VName, VName)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first Param Type -> VName
forall dec. Param dec -> VName
paramName) ([(Param Type, VName)] -> [(VName, VName)])
-> [(Param Type, VName)] -> [(VName, VName)]
forall a b. (a -> b) -> a -> b
$
        LoopNesting -> [(Param Type, VName)]
loopNestingParamsAndArrs LoopNesting
outermost

    identityStms :: (PatternT Type, [SubExp]) -> Stms rep
identityStms (PatternT Type
rem_pat, [SubExp]
res) =
      [Stm rep] -> Stms rep
forall rep. [Stm rep] -> Stms rep
stmsFromList ([Stm rep] -> Stms rep) -> [Stm rep] -> Stms rep
forall a b. (a -> b) -> a -> b
$ (PatElemT Type -> SubExp -> Stm rep)
-> [PatElemT Type] -> [SubExp] -> [Stm rep]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatElemT Type -> SubExp -> Stm rep
identityStm (PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements PatternT Type
rem_pat) [SubExp]
res
    identityStm :: PatElemT Type -> SubExp -> Stm rep
identityStm PatElemT Type
pe (Var VName
v)
      | Just VName
arr <- VName -> [(VName, VName)] -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v [(VName, VName)]
params_to_arrs =
        Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] [PatElemT Type
pe]) (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
arr
    identityStm PatElemT Type
pe SubExp
se =
      Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] [PatElemT Type
pe]) (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$
        BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$
          ShapeBase SubExp -> SubExp -> BasicOp
Replicate ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [LoopNesting -> SubExp
loopNestingWidth LoopNesting
outermost]) SubExp
se

addPostStms :: Monad m => PostStms rep -> DistNestT rep m ()
addPostStms :: forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
ks = DistRes rep -> DistNestT rep m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (DistRes rep -> DistNestT rep m ())
-> DistRes rep -> DistNestT rep m ()
forall a b. (a -> b) -> a -> b
$ DistRes Any
forall a. Monoid a => a
mempty {accPostStms :: PostStms rep
accPostStms = PostStms rep
ks}

postStm :: Monad m => Stms rep -> DistNestT rep m ()
postStm :: forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm Stms rep
stms = PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms (PostStms rep -> DistNestT rep m ())
-> PostStms rep -> DistNestT rep m ()
forall a b. (a -> b) -> a -> b
$ Stms rep -> PostStms rep
forall rep. Stms rep -> PostStms rep
PostStms Stms rep
stms

withStm ::
  (Monad m, DistRep rep) =>
  Stm SOACS ->
  DistNestT rep m a ->
  DistNestT rep m a
withStm :: forall (m :: * -> *) rep a.
(Monad m, DistRep rep) =>
Stm SOACS -> DistNestT rep m a -> DistNestT rep m a
withStm Stm SOACS
stm = (DistEnv rep m -> DistEnv rep m)
-> DistNestT rep m a -> DistNestT rep m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((DistEnv rep m -> DistEnv rep m)
 -> DistNestT rep m a -> DistNestT rep m a)
-> (DistEnv rep m -> DistEnv rep m)
-> DistNestT rep m a
-> DistNestT rep m a
forall a b. (a -> b) -> a -> b
$ \DistEnv rep m
env ->
  DistEnv rep m
env
    { distScope :: Scope rep
distScope =
        Scope SOACS -> Scope rep
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Stm SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm SOACS
stm) Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> DistEnv rep m -> Scope rep
forall rep (m :: * -> *). DistEnv rep m -> Scope rep
distScope DistEnv rep m
env,
      distNest :: Nestings
distNest =
        Names -> Nestings -> Nestings
letBindInInnerNesting Names
provided (Nestings -> Nestings) -> Nestings -> Nestings
forall a b. (a -> b) -> a -> b
$
          DistEnv rep m -> Nestings
forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest DistEnv rep m
env
    }
  where
    provided :: Names
provided = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT Type -> [VName]) -> PatternT Type -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Pattern
forall rep. Stm rep -> Pattern rep
stmPattern Stm SOACS
stm

leavingNesting ::
  (MonadFreshNames m, DistRep rep) =>
  DistAcc rep ->
  DistNestT rep m (DistAcc rep)
leavingNesting :: forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (DistAcc rep)
leavingNesting DistAcc rep
acc =
  case Targets -> Maybe ((PatternT Type, [SubExp]), Targets)
popInnerTarget (Targets -> Maybe ((PatternT Type, [SubExp]), Targets))
-> Targets -> Maybe ((PatternT Type, [SubExp]), Targets)
forall a b. (a -> b) -> a -> b
$ DistAcc rep -> Targets
forall rep. DistAcc rep -> Targets
distTargets DistAcc rep
acc of
    Maybe ((PatternT Type, [SubExp]), Targets)
Nothing ->
      [Char] -> DistNestT rep m (DistAcc rep)
forall a. HasCallStack => [Char] -> a
error [Char]
"The kernel targets list is unexpectedly small"
    Just ((PatternT Type
pat, [SubExp]
res), Targets
newtargets)
      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Seq (Stm rep) -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Seq (Stm rep) -> Bool) -> Seq (Stm rep) -> Bool
forall a b. (a -> b) -> a -> b
$ DistAcc rep -> Seq (Stm rep)
forall rep. DistAcc rep -> Stms rep
distStms DistAcc rep
acc -> do
        -- Any statements left over correspond to something that
        -- could not be distributed because it would cause irregular
        -- arrays.  These must be reconstructed into a a Map SOAC
        -- that will be sequentialised. XXX: life would be better if
        -- we were able to distribute irregular parallelism.
        (Nesting Names
_ LoopNesting
inner, [Nesting]
_) <- (DistEnv rep m -> Nestings) -> DistNestT rep m Nestings
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Nestings
forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest
        let MapNesting PatternT Type
_ StmAux ()
aux SubExp
w [(Param Type, VName)]
params_and_arrs = LoopNesting
inner
            body :: BodyT rep
body = BodyDec rep -> Seq (Stm rep) -> [SubExp] -> BodyT rep
forall rep. BodyDec rep -> Stms rep -> [SubExp] -> BodyT rep
Body () (DistAcc rep -> Seq (Stm rep)
forall rep. DistAcc rep -> Stms rep
distStms DistAcc rep
acc) [SubExp]
res
            used_in_body :: Names
used_in_body = BodyT rep -> Names
forall a. FreeIn a => a -> Names
freeIn BodyT rep
body
            ([Param Type]
used_params, [VName]
used_arrs) =
              [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Param Type, VName)] -> ([Param Type], [VName]))
-> [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. (a -> b) -> a -> b
$
                ((Param Type, VName) -> Bool)
-> [(Param Type, VName)] -> [(Param Type, VName)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Names -> Bool
`nameIn` Names
used_in_body) (VName -> Bool)
-> ((Param Type, VName) -> VName) -> (Param Type, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName)
-> ((Param Type, VName) -> Param Type)
-> (Param Type, VName)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type, VName) -> Param Type
forall a b. (a, b) -> a
fst) [(Param Type, VName)]
params_and_arrs
            lam' :: LambdaT rep
lam' =
              Lambda :: forall rep. [LParam rep] -> BodyT rep -> [Type] -> LambdaT rep
Lambda
                { lambdaParams :: [LParam rep]
lambdaParams = [Param Type]
[LParam rep]
used_params,
                  lambdaBody :: BodyT rep
lambdaBody = BodyT rep
body,
                  lambdaReturnType :: [Type]
lambdaReturnType = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ PatternT Type -> [Type]
forall dec. Typed dec => PatternT dec -> [Type]
patternTypes PatternT Type
pat
                }
        Seq (Stm rep)
stms <-
          Binder rep () -> DistNestT rep m (Seq (Stm rep))
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (Stms rep)
runBinder_ (Binder rep () -> DistNestT rep m (Seq (Stm rep)))
-> (SOAC rep -> Binder rep ())
-> SOAC rep
-> DistNestT rep m (Seq (Stm rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StmAux () -> Binder rep () -> Binder rep ()
forall (m :: * -> *) anyrep a.
MonadBinder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
aux (Binder rep () -> Binder rep ())
-> (SOAC rep -> Binder rep ()) -> SOAC rep -> Binder rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pattern (Rep (BinderT rep (State VNameSource)))
-> SOAC (Rep (BinderT rep (State VNameSource))) -> Binder rep ()
forall (m :: * -> *).
Transformer m =>
Pattern (Rep m) -> SOAC (Rep m) -> m ()
FOT.transformSOAC PatternT Type
Pattern (Rep (BinderT rep (State VNameSource)))
pat (SOAC rep -> DistNestT rep m (Seq (Stm rep)))
-> SOAC rep -> DistNestT rep m (Seq (Stm rep))
forall a b. (a -> b) -> a -> b
$
            SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
used_arrs (ScremaForm rep -> SOAC rep) -> ScremaForm rep -> SOAC rep
forall a b. (a -> b) -> a -> b
$ LambdaT rep -> ScremaForm rep
forall rep. Lambda rep -> ScremaForm rep
mapSOAC LambdaT rep
lam'

        DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ DistAcc rep
acc {distTargets :: Targets
distTargets = Targets
newtargets, distStms :: Seq (Stm rep)
distStms = Seq (Stm rep)
stms}
      | Bool
otherwise -> do
        -- Any results left over correspond to a Replicate or a Copy in
        -- the parent nesting, depending on whether the argument is a
        -- parameter of the innermost nesting.
        (Nesting Names
_ LoopNesting
inner_nesting, [Nesting]
_) <- (DistEnv rep m -> Nestings) -> DistNestT rep m Nestings
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Nestings
forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest
        let w :: SubExp
w = LoopNesting -> SubExp
loopNestingWidth LoopNesting
inner_nesting
            aux :: StmAux ()
aux = LoopNesting -> StmAux ()
loopNestingAux LoopNesting
inner_nesting
            inps :: [(Param Type, VName)]
inps = LoopNesting -> [(Param Type, VName)]
loopNestingParamsAndArrs LoopNesting
inner_nesting

            remnantStm :: PatElemT Type -> SubExp -> Stm rep
remnantStm PatElemT Type
pe (Var VName
v)
              | Just (Param Type
_, VName
arr) <- ((Param Type, VName) -> Bool)
-> [(Param Type, VName)] -> Maybe (Param Type, VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool)
-> ((Param Type, VName) -> VName) -> (Param Type, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName)
-> ((Param Type, VName) -> Param Type)
-> (Param Type, VName)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type, VName) -> Param Type
forall a b. (a, b) -> a
fst) [(Param Type, VName)]
inps =
                Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] [PatElemT Type
pe]) StmAux ()
StmAux (ExpDec rep)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
arr
            remnantStm PatElemT Type
pe SubExp
se =
              Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] [PatElemT Type
pe]) StmAux ()
StmAux (ExpDec rep)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$
                BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) SubExp
se

            stms :: Seq (Stm rep)
stms =
              [Stm rep] -> Seq (Stm rep)
forall rep. [Stm rep] -> Stms rep
stmsFromList ([Stm rep] -> Seq (Stm rep)) -> [Stm rep] -> Seq (Stm rep)
forall a b. (a -> b) -> a -> b
$ (PatElemT Type -> SubExp -> Stm rep)
-> [PatElemT Type] -> [SubExp] -> [Stm rep]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatElemT Type -> SubExp -> Stm rep
remnantStm (PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternElements PatternT Type
pat) [SubExp]
res

        DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ DistAcc rep
acc {distTargets :: Targets
distTargets = Targets
newtargets, distStms :: Seq (Stm rep)
distStms = Seq (Stm rep)
stms}

mapNesting ::
  (MonadFreshNames m, DistRep rep) =>
  PatternT Type ->
  StmAux () ->
  SubExp ->
  Lambda SOACS ->
  [VName] ->
  DistNestT rep m (DistAcc rep) ->
  DistNestT rep m (DistAcc rep)
mapNesting :: forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
PatternT Type
-> StmAux ()
-> SubExp
-> Lambda
-> [VName]
-> DistNestT rep m (DistAcc rep)
-> DistNestT rep m (DistAcc rep)
mapNesting PatternT Type
pat StmAux ()
aux SubExp
w Lambda
lam [VName]
arrs DistNestT rep m (DistAcc rep)
m =
  (DistEnv rep m -> DistEnv rep m)
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local DistEnv rep m -> DistEnv rep m
extend (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (DistAcc rep)
leavingNesting (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< DistNestT rep m (DistAcc rep)
m
  where
    nest :: Nesting
nest =
      Names -> LoopNesting -> Nesting
Nesting Names
forall a. Monoid a => a
mempty (LoopNesting -> Nesting) -> LoopNesting -> Nesting
forall a b. (a -> b) -> a -> b
$
        PatternT Type
-> StmAux () -> SubExp -> [(Param Type, VName)] -> LoopNesting
MapNesting PatternT Type
pat StmAux ()
aux SubExp
w ([(Param Type, VName)] -> LoopNesting)
-> [(Param Type, VName)] -> LoopNesting
forall a b. (a -> b) -> a -> b
$
          [Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda -> [LParam SOACS]
forall rep. LambdaT rep -> [LParam rep]
lambdaParams Lambda
lam) [VName]
arrs
    extend :: DistEnv rep m -> DistEnv rep m
extend DistEnv rep m
env =
      DistEnv rep m
env
        { distNest :: Nestings
distNest = Nesting -> Nestings -> Nestings
pushInnerNesting Nesting
nest (Nestings -> Nestings) -> Nestings -> Nestings
forall a b. (a -> b) -> a -> b
$ DistEnv rep m -> Nestings
forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest DistEnv rep m
env,
          distScope :: Scope rep
distScope = Scope SOACS -> Scope rep
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Lambda -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Lambda
lam) Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> DistEnv rep m -> Scope rep
forall rep (m :: * -> *). DistEnv rep m -> Scope rep
distScope DistEnv rep m
env
        }

inNesting ::
  (Monad m, DistRep rep) =>
  KernelNest ->
  DistNestT rep m a ->
  DistNestT rep m a
inNesting :: forall (m :: * -> *) rep a.
(Monad m, DistRep rep) =>
KernelNest -> DistNestT rep m a -> DistNestT rep m a
inNesting (LoopNesting
outer, [LoopNesting]
nests) = (DistEnv rep m -> DistEnv rep m)
-> DistNestT rep m a -> DistNestT rep m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((DistEnv rep m -> DistEnv rep m)
 -> DistNestT rep m a -> DistNestT rep m a)
-> (DistEnv rep m -> DistEnv rep m)
-> DistNestT rep m a
-> DistNestT rep m a
forall a b. (a -> b) -> a -> b
$ \DistEnv rep m
env ->
  DistEnv rep m
env
    { distNest :: Nestings
distNest = (Nesting
inner, [Nesting]
nests'),
      distScope :: Scope rep
distScope = (LoopNesting -> Scope rep) -> [LoopNesting] -> Scope rep
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap LoopNesting -> Scope rep
forall rep. DistRep rep => LoopNesting -> Scope rep
scopeOfLoopNesting (LoopNesting
outer LoopNesting -> [LoopNesting] -> [LoopNesting]
forall a. a -> [a] -> [a]
: [LoopNesting]
nests) Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> DistEnv rep m -> Scope rep
forall rep (m :: * -> *). DistEnv rep m -> Scope rep
distScope DistEnv rep m
env
    }
  where
    (Nesting
inner, [Nesting]
nests') =
      case [LoopNesting] -> [LoopNesting]
forall a. [a] -> [a]
reverse [LoopNesting]
nests of
        [] -> (LoopNesting -> Nesting
asNesting LoopNesting
outer, [])
        (LoopNesting
inner' : [LoopNesting]
ns) -> (LoopNesting -> Nesting
asNesting LoopNesting
inner', (LoopNesting -> Nesting) -> [LoopNesting] -> [Nesting]
forall a b. (a -> b) -> [a] -> [b]
map LoopNesting -> Nesting
asNesting ([LoopNesting] -> [Nesting]) -> [LoopNesting] -> [Nesting]
forall a b. (a -> b) -> a -> b
$ LoopNesting
outer LoopNesting -> [LoopNesting] -> [LoopNesting]
forall a. a -> [a] -> [a]
: [LoopNesting] -> [LoopNesting]
forall a. [a] -> [a]
reverse [LoopNesting]
ns)
    asNesting :: LoopNesting -> Nesting
asNesting = Names -> LoopNesting -> Nesting
Nesting Names
forall a. Monoid a => a
mempty

bodyContainsParallelism :: Body SOACS -> Bool
bodyContainsParallelism :: Body SOACS -> Bool
bodyContainsParallelism = (Stm SOACS -> Bool) -> Stms SOACS -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Stm SOACS -> Bool
isParallelStm (Stms SOACS -> Bool)
-> (Body SOACS -> Stms SOACS) -> Body SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body SOACS -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms
  where
    isParallelStm :: Stm SOACS -> Bool
isParallelStm Stm SOACS
stm =
      Exp SOACS -> Bool
isMap (Stm SOACS -> Exp SOACS
forall rep. Stm rep -> Exp rep
stmExp Stm SOACS
stm)
        Bool -> Bool -> Bool
&& 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))
    isMap :: Exp SOACS -> Bool
isMap BasicOp {} = Bool
False
    isMap Apply {} = Bool
False
    isMap If {} = Bool
False
    isMap (DoLoop [(FParam SOACS, SubExp)]
_ [(FParam SOACS, SubExp)]
_ ForLoop {} Body SOACS
body) = Body SOACS -> Bool
bodyContainsParallelism Body SOACS
body
    isMap (DoLoop [(FParam SOACS, SubExp)]
_ [(FParam SOACS, SubExp)]
_ WhileLoop {} Body SOACS
_) = Bool
False
    isMap (WithAcc [(ShapeBase SubExp, [VName], Maybe (Lambda, [SubExp]))]
_ Lambda
lam) = Body SOACS -> Bool
bodyContainsParallelism (Body SOACS -> Bool) -> Body SOACS -> Bool
forall a b. (a -> b) -> a -> b
$ Lambda -> Body SOACS
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda
lam
    isMap Op {} = Bool
True

lambdaContainsParallelism :: Lambda SOACS -> Bool
lambdaContainsParallelism :: Lambda -> Bool
lambdaContainsParallelism = Body SOACS -> Bool
bodyContainsParallelism (Body SOACS -> Bool) -> (Lambda -> Body SOACS) -> Lambda -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda -> Body SOACS
forall rep. LambdaT rep -> BodyT rep
lambdaBody

distributeMapBodyStms ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  DistAcc rep ->
  Stms SOACS ->
  DistNestT rep m (DistAcc rep)
distributeMapBodyStms :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
distributeMapBodyStms DistAcc rep
orig_acc = DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (DistAcc rep)
distribute (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> (Stms SOACS -> DistNestT rep m (DistAcc rep))
-> Stms SOACS
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< DistAcc rep -> [Stm SOACS] -> DistNestT rep m (DistAcc rep)
forall {rep} {m :: * -> *}.
(MonadFreshNames m, Bindable rep, HasSegOp rep, BinderOps rep,
 CanBeAliased (Op rep), LocalScope rep m, BodyDec rep ~ (),
 ExpDec rep ~ (), LetDec rep ~ Type) =>
DistAcc rep -> [Stm SOACS] -> DistNestT rep m (DistAcc rep)
onStms DistAcc rep
orig_acc ([Stm SOACS] -> DistNestT rep m (DistAcc rep))
-> (Stms SOACS -> [Stm SOACS])
-> Stms SOACS
-> DistNestT rep m (DistAcc rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList
  where
    onStms :: DistAcc rep -> [Stm SOACS] -> DistNestT rep m (DistAcc rep)
onStms DistAcc rep
acc [] = DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc
    onStms DistAcc rep
acc (Let Pattern
pat (StmAux Certificates
cs Attrs
_ ExpDec SOACS
_) (Op (Stream SubExp
w [VName]
arrs StreamForm SOACS
Sequential [SubExp]
accs Lambda
lam)) : [Stm SOACS]
stms) = do
      Scope SOACS
types <- (Scope rep -> Scope SOACS) -> DistNestT rep m (Scope SOACS)
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope rep -> Scope SOACS
forall rep. SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs
      Stms SOACS
stream_stms <-
        ((), Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> b
snd (((), Stms SOACS) -> Stms SOACS)
-> DistNestT rep m ((), Stms SOACS) -> DistNestT rep m (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BinderT SOACS (DistNestT rep m) ()
-> Scope SOACS -> DistNestT rep m ((), Stms SOACS)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BinderT rep m a -> Scope rep -> m (a, Stms rep)
runBinderT (Pattern (Rep (BinderT SOACS (DistNestT rep m)))
-> SubExp
-> [SubExp]
-> LambdaT (Rep (BinderT SOACS (DistNestT rep m)))
-> [VName]
-> BinderT SOACS (DistNestT rep m) ()
forall (m :: * -> *).
(MonadBinder m, Bindable (Rep m)) =>
Pattern (Rep m)
-> SubExp -> [SubExp] -> LambdaT (Rep m) -> [VName] -> m ()
sequentialStreamWholeArray Pattern (Rep (BinderT SOACS (DistNestT rep m)))
Pattern
pat SubExp
w [SubExp]
accs LambdaT (Rep (BinderT SOACS (DistNestT rep m)))
Lambda
lam [VName]
arrs) Scope SOACS
types
      (SymbolTable (Wise SOACS)
_, Stms SOACS
stream_stms') <-
        ReaderT
  (Scope SOACS)
  (DistNestT rep m)
  (SymbolTable (Wise SOACS), Stms SOACS)
-> Scope SOACS
-> DistNestT rep m (SymbolTable (Wise SOACS), Stms SOACS)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (SimpleOps SOACS
-> Scope SOACS
-> Stms SOACS
-> ReaderT
     (Scope SOACS)
     (DistNestT rep m)
     (SymbolTable (Wise SOACS), Stms SOACS)
forall (m :: * -> *) rep.
(MonadFreshNames m, SimplifiableRep rep) =>
SimpleOps rep
-> Scope rep -> Stms rep -> m (SymbolTable (Wise rep), Stms rep)
copyPropagateInStms SimpleOps SOACS
simpleSOACS Scope SOACS
types Stms SOACS
stream_stms) Scope SOACS
types
      DistAcc rep -> [Stm SOACS] -> DistNestT rep m (DistAcc rep)
onStms DistAcc rep
acc ([Stm SOACS] -> DistNestT rep m (DistAcc rep))
-> [Stm SOACS] -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList ((Stm SOACS -> Stm SOACS) -> Stms SOACS -> Stms SOACS
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certificates -> Stm SOACS -> Stm SOACS
forall rep. Certificates -> Stm rep -> Stm rep
certify Certificates
cs) Stms SOACS
stream_stms') [Stm SOACS] -> [Stm SOACS] -> [Stm SOACS]
forall a. [a] -> [a] -> [a]
++ [Stm SOACS]
stms
    onStms DistAcc rep
acc (Stm SOACS
stm : [Stm SOACS]
stms) =
      -- It is important that stm is in scope if 'maybeDistributeStm'
      -- wants to distribute, even if this causes the slightly silly
      -- situation that stm is in scope of itself.
      Stm SOACS
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep a.
(Monad m, DistRep rep) =>
Stm SOACS -> DistNestT rep m a -> DistNestT rep m a
withStm Stm SOACS
stm (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
maybeDistributeStm Stm SOACS
stm (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< DistAcc rep -> [Stm SOACS] -> DistNestT rep m (DistAcc rep)
onStms DistAcc rep
acc [Stm SOACS]
stms

onInnerMap :: Monad m => MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
onInnerMap :: forall (m :: * -> *) rep.
Monad m =>
MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
onInnerMap MapLoop
loop DistAcc rep
acc = do
  MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
f <- (DistEnv rep m
 -> MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistNestT
     rep m (MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m
-> MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *).
DistEnv rep m
-> MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
distOnInnerMap
  MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
f MapLoop
loop DistAcc rep
acc

onTopLevelStms :: Monad m => Stms SOACS -> DistNestT rep m ()
onTopLevelStms :: forall (m :: * -> *) rep.
Monad m =>
Stms SOACS -> DistNestT rep m ()
onTopLevelStms Stms SOACS
stms = do
  Stms SOACS -> DistNestT rep m (Stms rep)
f <- (DistEnv rep m -> Stms SOACS -> DistNestT rep m (Stms rep))
-> DistNestT rep m (Stms SOACS -> DistNestT rep m (Stms rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Stms SOACS -> DistNestT rep m (Stms rep)
forall rep (m :: * -> *).
DistEnv rep m -> Stms SOACS -> DistNestT rep m (Stms rep)
distOnTopLevelStms
  Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm (Stms rep -> DistNestT rep m ())
-> DistNestT rep m (Stms rep) -> DistNestT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Stms SOACS -> DistNestT rep m (Stms rep)
f Stms SOACS
stms

maybeDistributeStm ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  Stm SOACS ->
  DistAcc rep ->
  DistNestT rep m (DistAcc rep)
maybeDistributeStm :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
maybeDistributeStm Stm SOACS
stm DistAcc rep
acc
  | 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) =
    Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
maybeDistributeStm (Let Pattern
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac)) DistAcc rep
acc
  | Attr
"sequential_outer" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux =
    DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
distributeMapBodyStms DistAcc rep
acc (Stms SOACS -> DistNestT rep m (DistAcc rep))
-> (Stms SOACS -> Stms SOACS)
-> Stms SOACS
-> DistNestT rep m (DistAcc rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm SOACS -> Stm SOACS) -> Stms SOACS -> Stms SOACS
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certificates -> Stm SOACS -> Stm SOACS
forall rep. Certificates -> Stm rep -> Stm rep
certify (StmAux () -> Certificates
forall dec. StmAux dec -> Certificates
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux))
      (Stms SOACS -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (Stms SOACS) -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Binder SOACS () -> DistNestT rep m (Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (Stms rep)
runBinder_ (Pattern (Rep (BinderT SOACS (State VNameSource)))
-> SOAC (Rep (BinderT SOACS (State VNameSource)))
-> Binder SOACS ()
forall (m :: * -> *).
Transformer m =>
Pattern (Rep m) -> SOAC (Rep m) -> m ()
FOT.transformSOAC Pattern (Rep (BinderT SOACS (State VNameSource)))
Pattern
pat Op SOACS
SOAC (Rep (BinderT SOACS (State VNameSource)))
soac)
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
pat StmAux (ExpDec SOACS)
_ (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))) DistAcc rep
acc
  | Just Lambda
lam <- ScremaForm SOACS -> Maybe Lambda
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form =
    -- Only distribute inside the map if we can distribute everything
    -- following the map.
    DistAcc rep -> DistNestT rep m (Maybe (DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (Maybe (DistAcc rep))
distributeIfPossible DistAcc rep
acc DistNestT rep m (Maybe (DistAcc rep))
-> (Maybe (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Maybe (DistAcc rep)
Nothing -> Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
      Just DistAcc rep
acc' -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (DistAcc rep)
distribute (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
Monad m =>
MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
onInnerMap (Pattern -> StmAux () -> SubExp -> Lambda -> [VName] -> MapLoop
MapLoop Pattern
pat (Stm SOACS -> StmAux (ExpDec SOACS)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm SOACS
stm) SubExp
w Lambda
lam [VName]
arrs) DistAcc rep
acc'
maybeDistributeStm bnd :: Stm SOACS
bnd@(Let Pattern
pat StmAux (ExpDec SOACS)
aux (DoLoop [] [(FParam SOACS, SubExp)]
val form :: LoopForm SOACS
form@ForLoop {} Body SOACS
body)) DistAcc rep
acc
  | [PatElemT Type] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternContextElements PatternT Type
Pattern
pat),
    Body SOACS -> Bool
bodyContainsParallelism Body SOACS
body =
    DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
bnd DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
        | -- XXX: We cannot distribute if this loop depends on
          -- certificates bound within the loop nest (well, we could,
          -- but interchange would not be valid).  This is not a
          -- fundamental restriction, but an artifact of our
          -- certificate representation, which we should probably
          -- rethink.
          Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
            (LoopForm SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn LoopForm SOACS
form Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> StmAux () -> Names
forall a. FreeIn a => a -> Names
freeIn StmAux ()
StmAux (ExpDec SOACS)
aux)
              Names -> Names -> Bool
`namesIntersect` KernelNest -> Names
boundInKernelNest KernelNest
nest,
          Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res ->
          -- We need to pretend pat_unused was used anyway, by adding
          -- it to the kernel nest.
          Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
            PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
            KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest
            Scope SOACS
types <- (Scope rep -> Scope SOACS) -> DistNestT rep m (Scope SOACS)
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope rep -> Scope SOACS
forall rep. SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs

            -- Simplification is key to hoisting out statements that
            -- were variant to the loop, but invariant to the outer maps
            -- (which are now innermost).
            Stms SOACS
stms <-
              (ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> Scope SOACS -> DistNestT rep m (Stms SOACS)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
`runReaderT` Scope SOACS
types) (ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
 -> DistNestT rep m (Stms SOACS))
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> DistNestT rep m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$
                ((SymbolTable (Wise SOACS), Stms SOACS) -> Stms SOACS)
-> ReaderT
     (Scope SOACS)
     (DistNestT rep m)
     (SymbolTable (Wise SOACS), Stms SOACS)
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SymbolTable (Wise SOACS), Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> b
snd (ReaderT
   (Scope SOACS)
   (DistNestT rep m)
   (SymbolTable (Wise SOACS), Stms SOACS)
 -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS))
-> (Stms SOACS
    -> ReaderT
         (Scope SOACS)
         (DistNestT rep m)
         (SymbolTable (Wise SOACS), Stms SOACS))
-> Stms SOACS
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms SOACS
-> ReaderT
     (Scope SOACS)
     (DistNestT rep m)
     (SymbolTable (Wise SOACS), Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Stms SOACS -> m (SymbolTable (Wise SOACS), Stms SOACS)
simplifyStms
                  (Stms SOACS
 -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS))
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest
-> SeqLoop -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (m :: * -> *).
(MonadFreshNames m, HasScope SOACS m) =>
KernelNest -> SeqLoop -> m (Stms SOACS)
interchangeLoops KernelNest
nest' ([Int]
-> Pattern
-> [(FParam SOACS, SubExp)]
-> LoopForm SOACS
-> Body SOACS
-> SeqLoop
SeqLoop [Int]
perm Pattern
pat [(FParam SOACS, SubExp)]
val LoopForm SOACS
form Body SOACS
body)
            Stms SOACS -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
Stms SOACS -> DistNestT rep m ()
onTopLevelStms Stms SOACS
stms
            DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
      Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
        Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
bnd DistAcc rep
acc
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
pat StmAux (ExpDec SOACS)
_ (If SubExp
cond Body SOACS
tbranch Body SOACS
fbranch IfDec (BranchType SOACS)
ret)) DistAcc rep
acc
  | [PatElemT Type] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternContextElements PatternT Type
Pattern
pat),
    Body SOACS -> Bool
bodyContainsParallelism Body SOACS
tbranch Bool -> Bool -> Bool
|| Body SOACS -> Bool
bodyContainsParallelism Body SOACS
fbranch
      Bool -> Bool -> Bool
|| Bool -> Bool
not ((ExtType -> Bool) -> [ExtType] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ExtType -> Bool
forall shape u. TypeBase shape u -> Bool
primType (IfDec ExtType -> [ExtType]
forall rt. IfDec rt -> [rt]
ifReturns IfDec ExtType
IfDec (BranchType SOACS)
ret)) =
    DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
stm DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
            (SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
cond Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> IfDec ExtType -> Names
forall a. FreeIn a => a -> Names
freeIn IfDec ExtType
IfDec (BranchType SOACS)
ret) Names -> Names -> Bool
`namesIntersect` KernelNest -> Names
boundInKernelNest KernelNest
nest,
          Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res ->
          -- We need to pretend pat_unused was used anyway, by adding
          -- it to the kernel nest.
          Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
            KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest
            PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
            Scope SOACS
types <- (Scope rep -> Scope SOACS) -> DistNestT rep m (Scope SOACS)
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope rep -> Scope SOACS
forall rep. SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs
            let branch :: Branch
branch = [Int]
-> Pattern
-> SubExp
-> Body SOACS
-> Body SOACS
-> IfDec (BranchType SOACS)
-> Branch
Branch [Int]
perm Pattern
pat SubExp
cond Body SOACS
tbranch Body SOACS
fbranch IfDec (BranchType SOACS)
ret
            Stms SOACS
stms <-
              (ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> Scope SOACS -> DistNestT rep m (Stms SOACS)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
`runReaderT` Scope SOACS
types) (ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
 -> DistNestT rep m (Stms SOACS))
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> DistNestT rep m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$
                ((SymbolTable (Wise SOACS), Stms SOACS) -> Stms SOACS)
-> ReaderT
     (Scope SOACS)
     (DistNestT rep m)
     (SymbolTable (Wise SOACS), Stms SOACS)
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SymbolTable (Wise SOACS), Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> b
snd (ReaderT
   (Scope SOACS)
   (DistNestT rep m)
   (SymbolTable (Wise SOACS), Stms SOACS)
 -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS))
-> (Stms SOACS
    -> ReaderT
         (Scope SOACS)
         (DistNestT rep m)
         (SymbolTable (Wise SOACS), Stms SOACS))
-> Stms SOACS
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms SOACS
-> ReaderT
     (Scope SOACS)
     (DistNestT rep m)
     (SymbolTable (Wise SOACS), Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Stms SOACS -> m (SymbolTable (Wise SOACS), Stms SOACS)
simplifyStms
                  (Stms SOACS
 -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS))
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest
-> Branch -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (m :: * -> *).
(MonadFreshNames m, HasScope SOACS m) =>
KernelNest -> Branch -> m (Stms SOACS)
interchangeBranch KernelNest
nest' Branch
branch
            Stms SOACS -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
Stms SOACS -> DistNestT rep m ()
onTopLevelStms Stms SOACS
stms
            DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
      Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
        Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
pat StmAux (ExpDec SOACS)
_ (WithAcc [(ShapeBase SubExp, [VName], Maybe (Lambda, [SubExp]))]
inputs Lambda
lam)) DistAcc rep
acc
  | Lambda -> Bool
lambdaContainsParallelism Lambda
lam =
    DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
stm DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
            [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn (Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop Int
num_accs (Lambda -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda
lam))
              Names -> Names -> Bool
`namesIntersect` KernelNest -> Names
boundInKernelNest KernelNest
nest,
          Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res ->
          -- We need to pretend pat_unused was used anyway, by adding
          -- it to the kernel nest.
          Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
            KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest
            Scope SOACS
types <- (Scope rep -> Scope SOACS) -> DistNestT rep m (Scope SOACS)
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope rep -> Scope SOACS
forall rep. SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs
            PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
            let withacc :: WithAccStm
withacc = [Int]
-> Pattern
-> [(ShapeBase SubExp, [VName], Maybe (Lambda, [SubExp]))]
-> Lambda
-> WithAccStm
WithAccStm [Int]
perm Pattern
pat [(ShapeBase SubExp, [VName], Maybe (Lambda, [SubExp]))]
inputs Lambda
lam
            Stms SOACS
stms <-
              (ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> Scope SOACS -> DistNestT rep m (Stms SOACS)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
`runReaderT` Scope SOACS
types) (ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
 -> DistNestT rep m (Stms SOACS))
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> DistNestT rep m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$
                ((SymbolTable (Wise SOACS), Stms SOACS) -> Stms SOACS)
-> ReaderT
     (Scope SOACS)
     (DistNestT rep m)
     (SymbolTable (Wise SOACS), Stms SOACS)
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SymbolTable (Wise SOACS), Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> b
snd (ReaderT
   (Scope SOACS)
   (DistNestT rep m)
   (SymbolTable (Wise SOACS), Stms SOACS)
 -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS))
-> (Stms SOACS
    -> ReaderT
         (Scope SOACS)
         (DistNestT rep m)
         (SymbolTable (Wise SOACS), Stms SOACS))
-> Stms SOACS
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms SOACS
-> ReaderT
     (Scope SOACS)
     (DistNestT rep m)
     (SymbolTable (Wise SOACS), Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Stms SOACS -> m (SymbolTable (Wise SOACS), Stms SOACS)
simplifyStms (Stms SOACS
 -> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS))
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest
-> WithAccStm
-> ReaderT (Scope SOACS) (DistNestT rep m) (Stms SOACS)
forall (m :: * -> *).
(MonadFreshNames m, HasScope SOACS m) =>
KernelNest -> WithAccStm -> m (Stms SOACS)
interchangeWithAcc KernelNest
nest' WithAccStm
withacc
            Stms SOACS -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
Stms SOACS -> DistNestT rep m ()
onTopLevelStms Stms SOACS
stms
            DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
      Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
        Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
  where
    num_accs :: Int
num_accs = [(ShapeBase SubExp, [VName], Maybe (Lambda, [SubExp]))] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(ShapeBase SubExp, [VName], Maybe (Lambda, [SubExp]))]
inputs
maybeDistributeStm (Let Pattern
pat StmAux (ExpDec SOACS)
aux (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))) DistAcc rep
acc
  | Just [Reduce Commutativity
comm Lambda
lam [SubExp]
nes] <- ScremaForm SOACS -> Maybe [Reduce SOACS]
forall rep. ScremaForm rep -> Maybe [Reduce rep]
isReduceSOAC ScremaForm SOACS
form,
    Just BinderT SOACS (DistNestT rep m) ()
m <- Pattern
-> SubExp
-> Commutativity
-> Lambda
-> [(SubExp, VName)]
-> Maybe (BinderT SOACS (DistNestT rep m) ())
forall (m :: * -> *).
(MonadBinder m, Rep m ~ SOACS) =>
Pattern
-> SubExp
-> Commutativity
-> Lambda
-> [(SubExp, VName)]
-> Maybe (m ())
irwim Pattern
pat SubExp
w Commutativity
comm Lambda
lam ([(SubExp, VName)] -> Maybe (BinderT SOACS (DistNestT rep m) ()))
-> [(SubExp, VName)] -> Maybe (BinderT SOACS (DistNestT rep m) ())
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [VName]
arrs = do
    Scope SOACS
types <- (Scope rep -> Scope SOACS) -> DistNestT rep m (Scope SOACS)
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope rep -> Scope SOACS
forall rep. SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs
    (()
_, Stms SOACS
bnds) <- BinderT SOACS (DistNestT rep m) ()
-> Scope SOACS -> DistNestT rep m ((), Stms SOACS)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BinderT rep m a -> Scope rep -> m (a, Stms rep)
runBinderT (StmAux ()
-> BinderT SOACS (DistNestT rep m) ()
-> BinderT SOACS (DistNestT rep m) ()
forall (m :: * -> *) anyrep a.
MonadBinder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
StmAux (ExpDec SOACS)
aux BinderT SOACS (DistNestT rep m) ()
m) Scope SOACS
types
    DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
distributeMapBodyStms DistAcc rep
acc Stms SOACS
bnds

-- Parallelise segmented scatters.
maybeDistributeStm bnd :: Stm SOACS
bnd@(Let Pattern
pat (StmAux Certificates
cs Attrs
_ ExpDec SOACS
_) (Op (Scatter SubExp
w Lambda
lam [VName]
ivs [(ShapeBase SubExp, Int, VName)]
as))) DistAcc rep
acc =
  DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
bnd DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
      | Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res ->
        Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
          KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest
          Lambda rep
lam' <- Lambda -> DistNestT rep m (Lambda rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Lambda -> DistNestT rep m (Lambda rep)
soacsLambda Lambda
lam
          PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
          Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm (Stms rep -> DistNestT rep m ())
-> DistNestT rep m (Stms rep) -> DistNestT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest
-> [Int]
-> PatternT Type
-> Certificates
-> SubExp
-> Lambda rep
-> [VName]
-> [(ShapeBase SubExp, Int, VName)]
-> DistNestT rep m (Stms rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> PatternT Type
-> Certificates
-> SubExp
-> Lambda rep
-> [VName]
-> [(ShapeBase SubExp, Int, VName)]
-> DistNestT rep m (Stms rep)
segmentedScatterKernel KernelNest
nest' [Int]
perm PatternT Type
Pattern
pat Certificates
cs SubExp
w Lambda rep
lam' [VName]
ivs [(ShapeBase SubExp, Int, VName)]
as
          DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
    Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
      Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
bnd DistAcc rep
acc
-- Parallelise segmented Hist.
maybeDistributeStm bnd :: Stm SOACS
bnd@(Let Pattern
pat (StmAux Certificates
cs Attrs
_ ExpDec SOACS
_) (Op (Hist SubExp
w [HistOp SOACS]
ops Lambda
lam [VName]
as))) DistAcc rep
acc =
  DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
bnd DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
      | Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res ->
        Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
          Lambda rep
lam' <- Lambda -> DistNestT rep m (Lambda rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Lambda -> DistNestT rep m (Lambda rep)
soacsLambda Lambda
lam
          KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest
          PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
          Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm (Stms rep -> DistNestT rep m ())
-> DistNestT rep m (Stms rep) -> DistNestT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest
-> [Int]
-> Certificates
-> SubExp
-> [HistOp SOACS]
-> Lambda rep
-> [VName]
-> DistNestT rep m (Stms rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Certificates
-> SubExp
-> [HistOp SOACS]
-> Lambda rep
-> [VName]
-> DistNestT rep m (Stms rep)
segmentedHistKernel KernelNest
nest' [Int]
perm Certificates
cs SubExp
w [HistOp SOACS]
ops Lambda rep
lam' [VName]
as
          DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
    Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
      Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
bnd DistAcc rep
acc
-- Parallelise Index slices if the result is going to be returned
-- directly from the kernel.  This is because we would otherwise have
-- to sequentialise writing the result, which may be costly.
maybeDistributeStm
  stm :: Stm SOACS
stm@( Let
          (Pattern [] [PatElemT (LetDec SOACS)
pe])
          StmAux (ExpDec SOACS)
aux
          (BasicOp (Index VName
arr Slice SubExp
slice))
        )
  DistAcc rep
acc
    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([SubExp] -> Bool) -> [SubExp] -> Bool
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice,
      VName -> SubExp
Var (PatElemT Type -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT Type
PatElemT (LetDec SOACS)
pe) SubExp -> [SubExp] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` (PatternT Type, [SubExp]) -> [SubExp]
forall a b. (a, b) -> b
snd (Targets -> (PatternT Type, [SubExp])
innerTarget (DistAcc rep -> Targets
forall rep. DistAcc rep -> Targets
distTargets DistAcc rep
acc)) =
      DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
stm DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just (PostStms rep
kernels, [SubExp]
_res, KernelNest
nest, DistAcc rep
acc') ->
          Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
            PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
            Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm (Stms rep -> DistNestT rep m ())
-> DistNestT rep m (Stms rep) -> DistNestT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest
-> Certificates
-> VName
-> Slice SubExp
-> DistNestT rep m (Stms rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> Certificates
-> VName
-> Slice SubExp
-> DistNestT rep m (Stms rep)
segmentedGatherKernel KernelNest
nest (StmAux () -> Certificates
forall dec. StmAux dec -> Certificates
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux) VName
arr Slice SubExp
slice
            DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
        Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
          Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
-- If the scan can be distributed by itself, we will turn it into a
-- segmented scan.
--
-- If the scan cannot be distributed by itself, it will be
-- sequentialised in the default case for this function.
maybeDistributeStm bnd :: Stm SOACS
bnd@(Let Pattern
pat (StmAux Certificates
cs Attrs
_ ExpDec SOACS
_) (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))) DistAcc rep
acc
  | Just ([Scan SOACS]
scans, Lambda
map_lam) <- ScremaForm SOACS -> Maybe ([Scan SOACS], Lambda)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
isScanomapSOAC ScremaForm SOACS
form,
    Scan Lambda
lam [SubExp]
nes <- [Scan SOACS] -> Scan SOACS
forall rep. Bindable rep => [Scan rep] -> Scan rep
singleScan [Scan SOACS]
scans =
    DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
bnd DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
        | Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res ->
          -- We need to pretend pat_unused was used anyway, by adding
          -- it to the kernel nest.
          Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
            KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest
            Lambda rep
map_lam' <- Lambda -> DistNestT rep m (Lambda rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Lambda -> DistNestT rep m (Lambda rep)
soacsLambda Lambda
map_lam
            Lambda rep
lam' <- Lambda -> DistNestT rep m (Lambda rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Lambda -> DistNestT rep m (Lambda rep)
soacsLambda Lambda
lam
            Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$
              KernelNest
-> [Int]
-> SubExp
-> Lambda rep
-> Lambda rep
-> [SubExp]
-> [VName]
-> DistNestT rep m (Maybe (Stms rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> SubExp
-> Lambda rep
-> Lambda rep
-> [SubExp]
-> [VName]
-> DistNestT rep m (Maybe (Stms rep))
segmentedScanomapKernel KernelNest
nest' [Int]
perm SubExp
w Lambda rep
lam' Lambda rep
map_lam' [SubExp]
nes [VName]
arrs
                DistNestT rep m (Maybe (Stms rep))
-> (Maybe (Stms rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Certificates
-> Stm SOACS
-> DistAcc rep
-> PostStms rep
-> DistAcc rep
-> Maybe (Stms rep)
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Certificates
-> Stm SOACS
-> DistAcc rep
-> PostStms rep
-> DistAcc rep
-> Maybe (Stms rep)
-> DistNestT rep m (DistAcc rep)
kernelOrNot Certificates
cs Stm SOACS
bnd DistAcc rep
acc PostStms rep
kernels DistAcc rep
acc'
      Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
        Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
bnd DistAcc rep
acc
-- if the reduction can be distributed by itself, we will turn it into a
-- segmented reduce.
--
-- If the reduction cannot be distributed by itself, it will be
-- sequentialised in the default case for this function.
maybeDistributeStm bnd :: Stm SOACS
bnd@(Let Pattern
pat (StmAux Certificates
cs Attrs
_ ExpDec SOACS
_) (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))) DistAcc rep
acc
  | Just ([Reduce SOACS]
reds, Lambda
map_lam) <- ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm SOACS
form,
    Reduce Commutativity
comm Lambda
lam [SubExp]
nes <- [Reduce SOACS] -> Reduce SOACS
forall rep. Bindable rep => [Reduce rep] -> Reduce rep
singleReduce [Reduce SOACS]
reds =
    DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
bnd DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
        | Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res ->
          -- We need to pretend pat_unused was used anyway, by adding
          -- it to the kernel nest.
          Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
            KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest

            Lambda rep
lam' <- Lambda -> DistNestT rep m (Lambda rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Lambda -> DistNestT rep m (Lambda rep)
soacsLambda Lambda
lam
            Lambda rep
map_lam' <- Lambda -> DistNestT rep m (Lambda rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Lambda -> DistNestT rep m (Lambda rep)
soacsLambda Lambda
map_lam

            let comm' :: Commutativity
comm'
                  | Lambda -> Bool
forall rep. Lambda rep -> Bool
commutativeLambda Lambda
lam = Commutativity
Commutative
                  | Bool
otherwise = Commutativity
comm

            KernelNest
-> [Int]
-> SubExp
-> Commutativity
-> Lambda rep
-> Lambda rep
-> [SubExp]
-> [VName]
-> DistNestT rep m (Maybe (Stms rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> SubExp
-> Commutativity
-> Lambda rep
-> Lambda rep
-> [SubExp]
-> [VName]
-> DistNestT rep m (Maybe (Stms rep))
regularSegmentedRedomapKernel KernelNest
nest' [Int]
perm SubExp
w Commutativity
comm' Lambda rep
lam' Lambda rep
map_lam' [SubExp]
nes [VName]
arrs
              DistNestT rep m (Maybe (Stms rep))
-> (Maybe (Stms rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Certificates
-> Stm SOACS
-> DistAcc rep
-> PostStms rep
-> DistAcc rep
-> Maybe (Stms rep)
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Certificates
-> Stm SOACS
-> DistAcc rep
-> PostStms rep
-> DistAcc rep
-> Maybe (Stms rep)
-> DistNestT rep m (DistAcc rep)
kernelOrNot Certificates
cs Stm SOACS
bnd DistAcc rep
acc PostStms rep
kernels DistAcc rep
acc'
      Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
        Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
bnd DistAcc rep
acc
maybeDistributeStm (Let Pattern
pat (StmAux Certificates
cs Attrs
_ ExpDec SOACS
_) (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))) DistAcc rep
acc = do
  -- This Screma is too complicated for us to immediately do
  -- anything, so split it up and try again.
  Scope SOACS
scope <- (Scope rep -> Scope SOACS) -> DistNestT rep m (Scope SOACS)
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope rep -> Scope SOACS
forall rep. SameScope rep SOACS => Scope rep -> Scope SOACS
scopeForSOACs
  DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
distributeMapBodyStms DistAcc rep
acc (Stms SOACS -> DistNestT rep m (DistAcc rep))
-> (((), Stms SOACS) -> Stms SOACS)
-> ((), Stms SOACS)
-> DistNestT rep m (DistAcc rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm SOACS -> Stm SOACS) -> Stms SOACS -> Stms SOACS
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certificates -> Stm SOACS -> Stm SOACS
forall rep. Certificates -> Stm rep -> Stm rep
certify Certificates
cs) (Stms SOACS -> Stms SOACS)
-> (((), Stms SOACS) -> Stms SOACS)
-> ((), Stms SOACS)
-> Stms SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((), Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> b
snd
    (((), Stms SOACS) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m ((), Stms SOACS)
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinderT SOACS (DistNestT rep m) ()
-> Scope SOACS -> DistNestT rep m ((), Stms SOACS)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BinderT rep m a -> Scope rep -> m (a, Stms rep)
runBinderT (Pattern (Rep (BinderT SOACS (DistNestT rep m)))
-> SubExp
-> ScremaForm (Rep (BinderT SOACS (DistNestT rep m)))
-> [VName]
-> BinderT SOACS (DistNestT rep m) ()
forall (m :: * -> *).
(MonadBinder m, Op (Rep m) ~ SOAC (Rep m), Bindable (Rep m)) =>
Pattern (Rep m) -> SubExp -> ScremaForm (Rep m) -> [VName] -> m ()
dissectScrema Pattern (Rep (BinderT SOACS (DistNestT rep m)))
Pattern
pat SubExp
w ScremaForm (Rep (BinderT SOACS (DistNestT rep m)))
ScremaForm SOACS
form [VName]
arrs) Scope SOACS
scope
maybeDistributeStm (Let Pattern
pat StmAux (ExpDec SOACS)
aux (BasicOp (Replicate (Shape (SubExp
d : [SubExp]
ds)) SubExp
v))) DistAcc rep
acc
  | [Type
t] <- PatternT Type -> [Type]
forall dec. Typed dec => PatternT dec -> [Type]
patternTypes PatternT Type
Pattern
pat = do
    VName
tmp <- [Char] -> DistNestT rep m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"tmp"
    let rowt :: Type
rowt = Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType Type
t
        newbnd :: Stm SOACS
newbnd = Pattern -> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pattern
pat StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ Op SOACS -> Exp SOACS
forall rep. Op rep -> ExpT rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
d [] (ScremaForm SOACS -> SOAC SOACS) -> ScremaForm SOACS -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ Lambda -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda
lam
        tmpbnd :: Stm SOACS
tmpbnd =
          Pattern -> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] [VName -> Type -> PatElemT Type
forall dec. VName -> dec -> PatElemT dec
PatElem VName
tmp Type
rowt]) StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
            BasicOp -> Exp SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp]
ds) SubExp
v
        lam :: Lambda
lam =
          Lambda :: forall rep. [LParam rep] -> BodyT rep -> [Type] -> LambdaT rep
Lambda
            { lambdaReturnType :: [Type]
lambdaReturnType = [Type
rowt],
              lambdaParams :: [LParam SOACS]
lambdaParams = [],
              lambdaBody :: Body SOACS
lambdaBody = Stms SOACS -> [SubExp] -> Body SOACS
forall rep. Bindable rep => Stms rep -> [SubExp] -> Body rep
mkBody (Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
tmpbnd) [VName -> SubExp
Var VName
tmp]
            }
    Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
maybeDistributeStm Stm SOACS
newbnd DistAcc rep
acc
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
_ StmAux (ExpDec SOACS)
aux (BasicOp (Copy VName
stm_arr))) DistAcc rep
acc =
  DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
distributeSingleUnaryStm DistAcc rep
acc Stm SOACS
stm VName
stm_arr ((KernelNest
  -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
 -> DistNestT rep m (DistAcc rep))
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ \KernelNest
_ PatternT Type
outerpat VName
arr ->
    Stms rep -> DistNestT rep m (Stms rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms rep -> DistNestT rep m (Stms rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$ Stm rep -> Stms rep
forall rep. Stm rep -> Stms rep
oneStm (Stm rep -> Stms rep) -> Stm rep -> Stms rep
forall a b. (a -> b) -> a -> b
$ Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let PatternT Type
Pattern rep
outerpat StmAux (ExpDec rep)
StmAux (ExpDec SOACS)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
arr
-- Opaques are applied to the full array, because otherwise they can
-- drastically inhibit parallelisation in some cases.
maybeDistributeStm stm :: Stm SOACS
stm@(Let (Pattern [] [PatElemT (LetDec SOACS)
pe]) StmAux (ExpDec SOACS)
aux (BasicOp (Opaque (Var VName
stm_arr)))) DistAcc rep
acc
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ PatElemT Type -> Type
forall t. Typed t => t -> Type
typeOf PatElemT Type
PatElemT (LetDec SOACS)
pe =
    DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
distributeSingleUnaryStm DistAcc rep
acc Stm SOACS
stm VName
stm_arr ((KernelNest
  -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
 -> DistNestT rep m (DistAcc rep))
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ \KernelNest
_ PatternT Type
outerpat VName
arr ->
      Stms rep -> DistNestT rep m (Stms rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms rep -> DistNestT rep m (Stms rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$ Stm rep -> Stms rep
forall rep. Stm rep -> Stms rep
oneStm (Stm rep -> Stms rep) -> Stm rep -> Stms rep
forall a b. (a -> b) -> a -> b
$ Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let PatternT Type
Pattern rep
outerpat StmAux (ExpDec rep)
StmAux (ExpDec SOACS)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
arr
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
_ StmAux (ExpDec SOACS)
aux (BasicOp (Rearrange [Int]
perm VName
stm_arr))) DistAcc rep
acc =
  DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
distributeSingleUnaryStm DistAcc rep
acc Stm SOACS
stm VName
stm_arr ((KernelNest
  -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
 -> DistNestT rep m (DistAcc rep))
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ \KernelNest
nest PatternT Type
outerpat VName
arr -> do
    let r :: Int
r = [LoopNesting] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (KernelNest -> [LoopNesting]
forall a b. (a, b) -> b
snd KernelNest
nest) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        perm' :: [Int]
perm' = [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) [Int]
perm
    -- We need to add a copy, because the original map nest
    -- will have produced an array without aliases, and so must we.
    VName
arr' <- [Char] -> DistNestT rep m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> DistNestT rep m VName)
-> [Char] -> DistNestT rep m VName
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString VName
arr
    Type
arr_t <- VName -> DistNestT rep m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
    Stms rep -> DistNestT rep m (Stms rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms rep -> DistNestT rep m (Stms rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$
      [Stm rep] -> Stms rep
forall rep. [Stm rep] -> Stms rep
stmsFromList
        [ Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] [VName -> Type -> PatElemT Type
forall dec. VName -> dec -> PatElemT dec
PatElem VName
arr' Type
arr_t]) StmAux (ExpDec rep)
StmAux (ExpDec SOACS)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
arr,
          Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let PatternT Type
Pattern rep
outerpat StmAux (ExpDec rep)
StmAux (ExpDec SOACS)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
Rearrange [Int]
perm' VName
arr'
        ]
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
_ StmAux (ExpDec SOACS)
aux (BasicOp (Reshape ShapeChange SubExp
reshape VName
stm_arr))) DistAcc rep
acc =
  DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
distributeSingleUnaryStm DistAcc rep
acc Stm SOACS
stm VName
stm_arr ((KernelNest
  -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
 -> DistNestT rep m (DistAcc rep))
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ \KernelNest
nest PatternT Type
outerpat VName
arr -> do
    let reshape' :: ShapeChange SubExp
reshape' =
          (SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew (KernelNest -> [SubExp]
kernelNestWidths KernelNest
nest)
            ShapeChange SubExp -> ShapeChange SubExp -> ShapeChange SubExp
forall a. [a] -> [a] -> [a]
++ (SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew (ShapeChange SubExp -> [SubExp]
forall d. ShapeChange d -> [d]
newDims ShapeChange SubExp
reshape)
    Stms rep -> DistNestT rep m (Stms rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms rep -> DistNestT rep m (Stms rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$ Stm rep -> Stms rep
forall rep. Stm rep -> Stms rep
oneStm (Stm rep -> Stms rep) -> Stm rep -> Stms rep
forall a b. (a -> b) -> a -> b
$ Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let PatternT Type
Pattern rep
outerpat StmAux (ExpDec rep)
StmAux (ExpDec SOACS)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
Reshape ShapeChange SubExp
reshape' VName
arr
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
_ StmAux (ExpDec SOACS)
aux (BasicOp (Rotate [SubExp]
rots VName
stm_arr))) DistAcc rep
acc =
  DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
distributeSingleUnaryStm DistAcc rep
acc Stm SOACS
stm VName
stm_arr ((KernelNest
  -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
 -> DistNestT rep m (DistAcc rep))
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ \KernelNest
nest PatternT Type
outerpat VName
arr -> do
    let rots' :: [SubExp]
rots' = (SubExp -> SubExp) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> SubExp -> SubExp
forall a b. a -> b -> a
const (SubExp -> SubExp -> SubExp) -> SubExp -> SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) (KernelNest -> [SubExp]
kernelNestWidths KernelNest
nest) [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
rots
    Stms rep -> DistNestT rep m (Stms rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms rep -> DistNestT rep m (Stms rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$ Stm rep -> Stms rep
forall rep. Stm rep -> Stms rep
oneStm (Stm rep -> Stms rep) -> Stm rep -> Stms rep
forall a b. (a -> b) -> a -> b
$ Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let PatternT Type
Pattern rep
outerpat StmAux (ExpDec rep)
StmAux (ExpDec SOACS)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> BasicOp
Rotate [SubExp]
rots' VName
arr
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
pat StmAux (ExpDec SOACS)
aux (BasicOp (Update VName
arr Slice SubExp
slice (Var VName
v)))) DistAcc rep
acc
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([SubExp] -> Bool) -> [SubExp] -> Bool
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice =
    DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
stm DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
        | [SubExp]
res [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT Type -> [VName]) -> PatternT Type -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Pattern
forall rep. Stm rep -> Pattern rep
stmPattern Stm SOACS
stm),
          Just ([Int]
perm, [PatElemT Type]
pat_unused) <- PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
Pattern
pat [SubExp]
res -> do
          PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
          Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
            KernelNest
nest' <- [PatElemT Type] -> KernelNest -> DistNestT rep m KernelNest
forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pat_unused KernelNest
nest
            Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm
              (Stms rep -> DistNestT rep m ())
-> DistNestT rep m (Stms rep) -> DistNestT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest
-> [Int]
-> Certificates
-> VName
-> Slice SubExp
-> VName
-> DistNestT rep m (Stms rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Certificates
-> VName
-> Slice SubExp
-> VName
-> DistNestT rep m (Stms rep)
segmentedUpdateKernel KernelNest
nest' [Int]
perm (StmAux () -> Certificates
forall dec. StmAux dec -> Certificates
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux) VName
arr Slice SubExp
slice VName
v
            DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
      Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ -> Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
maybeDistributeStm stm :: Stm SOACS
stm@(Let Pattern
_ StmAux (ExpDec SOACS)
aux (BasicOp (Concat Int
d VName
x [VName]
xs SubExp
w))) DistAcc rep
acc =
  DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
stm DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (PostStms rep
kernels, [SubExp]
_, KernelNest
nest, DistAcc rep
acc') ->
      Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$
        KernelNest -> DistNestT rep m (Maybe (Stms rep))
segmentedConcat KernelNest
nest
          DistNestT rep m (Maybe (Stms rep))
-> (Maybe (Stms rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Certificates
-> Stm SOACS
-> DistAcc rep
-> PostStms rep
-> DistAcc rep
-> Maybe (Stms rep)
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Certificates
-> Stm SOACS
-> DistAcc rep
-> PostStms rep
-> DistAcc rep
-> Maybe (Stms rep)
-> DistNestT rep m (DistAcc rep)
kernelOrNot (StmAux () -> Certificates
forall dec. StmAux dec -> Certificates
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux) Stm SOACS
stm DistAcc rep
acc PostStms rep
kernels DistAcc rep
acc'
    Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ ->
      Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
  where
    segmentedConcat :: KernelNest -> DistNestT rep m (Maybe (Stms rep))
segmentedConcat KernelNest
nest =
      KernelNest
-> [Int]
-> Names
-> Names
-> [SubExp]
-> [VName]
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Names
-> Names
-> [SubExp]
-> [VName]
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
isSegmentedOp KernelNest
nest [Int
0] Names
forall a. Monoid a => a
mempty Names
forall a. Monoid a => a
mempty [] (VName
x VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
xs) ((PatternT Type
  -> [(VName, SubExp)]
  -> [KernelInput]
  -> [SubExp]
  -> [VName]
  -> BinderT rep m ())
 -> DistNestT rep m (Maybe (Stms rep)))
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
forall a b. (a -> b) -> a -> b
$
        \PatternT Type
pat [(VName, SubExp)]
_ [KernelInput]
_ [SubExp]
_ (VName
x' : [VName]
xs') ->
          let d' :: Int
d' = Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [LoopNesting] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (KernelNest -> [LoopNesting]
forall a b. (a, b) -> b
snd KernelNest
nest) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
           in Stm (Rep (BinderT rep m)) -> BinderT rep m ()
forall (m :: * -> *). MonadBinder m => Stm (Rep m) -> m ()
addStm (Stm (Rep (BinderT rep m)) -> BinderT rep m ())
-> Stm (Rep (BinderT rep m)) -> BinderT rep m ()
forall a b. (a -> b) -> a -> b
$ Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let PatternT Type
Pattern rep
pat StmAux (ExpDec rep)
StmAux (ExpDec SOACS)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ Int -> VName -> [VName] -> SubExp -> BasicOp
Concat Int
d' VName
x' [VName]
xs' SubExp
w
maybeDistributeStm Stm SOACS
bnd DistAcc rep
acc =
  Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
bnd DistAcc rep
acc

distributeSingleUnaryStm ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  DistAcc rep ->
  Stm SOACS ->
  VName ->
  (KernelNest -> PatternT Type -> VName -> DistNestT rep m (Stms rep)) ->
  DistNestT rep m (DistAcc rep)
distributeSingleUnaryStm :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> VName
-> (KernelNest
    -> PatternT Type -> VName -> DistNestT rep m (Stms rep))
-> DistNestT rep m (DistAcc rep)
distributeSingleUnaryStm DistAcc rep
acc Stm SOACS
stm VName
stm_arr KernelNest -> PatternT Type -> VName -> DistNestT rep m (Stms rep)
f =
  DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
stm DistNestT
  rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
-> (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
    -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (PostStms rep
kernels, [SubExp]
res, KernelNest
nest, DistAcc rep
acc')
      | [SubExp]
res [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT Type -> [VName]) -> PatternT Type -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Pattern
forall rep. Stm rep -> Pattern rep
stmPattern Stm SOACS
stm),
        (LoopNesting
outer, [LoopNesting]
_) <- KernelNest
nest,
        [(Param Type
arr_p, VName
arr)] <- LoopNesting -> [(Param Type, VName)]
loopNestingParamsAndArrs LoopNesting
outer,
        KernelNest -> Names
boundInKernelNest KernelNest
nest Names -> Names -> Names
`namesIntersection` Stm SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Stm SOACS
stm
          Names -> Names -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> Names
oneName (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
arr_p),
        VName -> KernelNest -> Bool
perfectlyMapped VName
arr KernelNest
nest -> do
        PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
        let outerpat :: PatternT Type
outerpat = LoopNesting -> PatternT Type
loopNestingPattern (LoopNesting -> PatternT Type) -> LoopNesting -> PatternT Type
forall a b. (a -> b) -> a -> b
$ KernelNest -> LoopNesting
forall a b. (a, b) -> a
fst KernelNest
nest
        Scope rep
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (DistAcc rep -> Scope rep
forall rep. DistRep rep => DistAcc rep -> Scope rep
typeEnvFromDistAcc DistAcc rep
acc') (DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall a b. (a -> b) -> a -> b
$ do
          Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm (Stms rep -> DistNestT rep m ())
-> DistNestT rep m (Stms rep) -> DistNestT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< KernelNest -> PatternT Type -> VName -> DistNestT rep m (Stms rep)
f KernelNest
nest PatternT Type
outerpat VName
arr
          DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'
    Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
_ -> Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc Stm SOACS
stm DistAcc rep
acc
  where
    perfectlyMapped :: VName -> KernelNest -> Bool
perfectlyMapped VName
arr (LoopNesting
outer, [LoopNesting]
nest)
      | [(Param Type
p, VName
arr')] <- LoopNesting -> [(Param Type, VName)]
loopNestingParamsAndArrs LoopNesting
outer,
        VName
arr VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
arr' =
        case [LoopNesting]
nest of
          [] -> Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
stm_arr
          LoopNesting
x : [LoopNesting]
xs -> VName -> KernelNest -> Bool
perfectlyMapped (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p) (LoopNesting
x, [LoopNesting]
xs)
      | Bool
otherwise =
        Bool
False

distribute ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  DistAcc rep ->
  DistNestT rep m (DistAcc rep)
distribute :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (DistAcc rep)
distribute DistAcc rep
acc =
  DistAcc rep -> Maybe (DistAcc rep) -> DistAcc rep
forall a. a -> Maybe a -> a
fromMaybe DistAcc rep
acc (Maybe (DistAcc rep) -> DistAcc rep)
-> DistNestT rep m (Maybe (DistAcc rep))
-> DistNestT rep m (DistAcc rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DistAcc rep -> DistNestT rep m (Maybe (DistAcc rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (Maybe (DistAcc rep))
distributeIfPossible DistAcc rep
acc

mkSegLevel ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  DistNestT rep m (MkSegLevel rep (DistNestT rep m))
mkSegLevel :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistNestT rep m (MkSegLevel rep (DistNestT rep m))
mkSegLevel = do
  [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl <- (DistEnv rep m
 -> [SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep m (SegOpLevel rep))
-> DistNestT
     rep
     m
     ([SubExp]
      -> [Char]
      -> ThreadRecommendation
      -> BinderT rep m (SegOpLevel rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m
-> [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep m (SegOpLevel rep)
forall rep (m :: * -> *). DistEnv rep m -> MkSegLevel rep m
distSegLevel
  MkSegLevel rep (DistNestT rep m)
-> DistNestT rep m (MkSegLevel rep (DistNestT rep m))
forall (m :: * -> *) a. Monad m => a -> m a
return (MkSegLevel rep (DistNestT rep m)
 -> DistNestT rep m (MkSegLevel rep (DistNestT rep m)))
-> MkSegLevel rep (DistNestT rep m)
-> DistNestT rep m (MkSegLevel rep (DistNestT rep m))
forall a b. (a -> b) -> a -> b
$ \[SubExp]
w [Char]
desc ThreadRecommendation
r -> do
    (SegOpLevel rep
lvl, Stms rep
stms) <- DistNestT rep m (SegOpLevel rep, Stms rep)
-> BinderT rep (DistNestT rep m) (SegOpLevel rep, Stms rep)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (DistNestT rep m (SegOpLevel rep, Stms rep)
 -> BinderT rep (DistNestT rep m) (SegOpLevel rep, Stms rep))
-> DistNestT rep m (SegOpLevel rep, Stms rep)
-> BinderT rep (DistNestT rep m) (SegOpLevel rep, Stms rep)
forall a b. (a -> b) -> a -> b
$ m (SegOpLevel rep, Stms rep)
-> DistNestT rep m (SegOpLevel rep, Stms rep)
forall rep (m :: * -> *) a.
(LocalScope rep m, DistRep rep) =>
m a -> DistNestT rep m a
liftInner (m (SegOpLevel rep, Stms rep)
 -> DistNestT rep m (SegOpLevel rep, Stms rep))
-> m (SegOpLevel rep, Stms rep)
-> DistNestT rep m (SegOpLevel rep, Stms rep)
forall a b. (a -> b) -> a -> b
$ BinderT rep m (SegOpLevel rep) -> m (SegOpLevel rep, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
BinderT rep m a -> m (a, Stms rep)
runBinderT' (BinderT rep m (SegOpLevel rep) -> m (SegOpLevel rep, Stms rep))
-> BinderT rep m (SegOpLevel rep) -> m (SegOpLevel rep, Stms rep)
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl [SubExp]
w [Char]
desc ThreadRecommendation
r
    Stms (Rep (BinderT rep (DistNestT rep m)))
-> BinderT rep (DistNestT rep m) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms rep
Stms (Rep (BinderT rep (DistNestT rep m)))
stms
    SegOpLevel rep -> BinderT rep (DistNestT rep m) (SegOpLevel rep)
forall (m :: * -> *) a. Monad m => a -> m a
return SegOpLevel rep
lvl

distributeIfPossible ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  DistAcc rep ->
  DistNestT rep m (Maybe (DistAcc rep))
distributeIfPossible :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (Maybe (DistAcc rep))
distributeIfPossible DistAcc rep
acc = do
  Nestings
nest <- (DistEnv rep m -> Nestings) -> DistNestT rep m Nestings
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Nestings
forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest
  [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl <- DistNestT
  rep
  m
  ([SubExp]
   -> [Char]
   -> ThreadRecommendation
   -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistNestT rep m (MkSegLevel rep (DistNestT rep m))
mkSegLevel
  ([SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
-> Nestings
-> Targets
-> Stms rep
-> DistNestT rep m (Maybe (Targets, Stms rep))
forall rep (m :: * -> *).
(DistRep rep, MonadFreshNames m, LocalScope rep m,
 MonadLogger m) =>
MkSegLevel rep m
-> Nestings -> Targets -> Stms rep -> m (Maybe (Targets, Stms rep))
tryDistribute [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl Nestings
nest (DistAcc rep -> Targets
forall rep. DistAcc rep -> Targets
distTargets DistAcc rep
acc) (DistAcc rep -> Stms rep
forall rep. DistAcc rep -> Stms rep
distStms DistAcc rep
acc) DistNestT rep m (Maybe (Targets, Stms rep))
-> (Maybe (Targets, Stms rep)
    -> DistNestT rep m (Maybe (DistAcc rep)))
-> DistNestT rep m (Maybe (DistAcc rep))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe (Targets, Stms rep)
Nothing -> Maybe (DistAcc rep) -> DistNestT rep m (Maybe (DistAcc rep))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (DistAcc rep)
forall a. Maybe a
Nothing
    Just (Targets
targets, Stms rep
kernel) -> do
      Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm Stms rep
kernel
      Maybe (DistAcc rep) -> DistNestT rep m (Maybe (DistAcc rep))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (DistAcc rep) -> DistNestT rep m (Maybe (DistAcc rep)))
-> Maybe (DistAcc rep) -> DistNestT rep m (Maybe (DistAcc rep))
forall a b. (a -> b) -> a -> b
$
        DistAcc rep -> Maybe (DistAcc rep)
forall a. a -> Maybe a
Just
          DistAcc :: forall rep. Targets -> Stms rep -> DistAcc rep
DistAcc
            { distTargets :: Targets
distTargets = Targets
targets,
              distStms :: Stms rep
distStms = Stms rep
forall a. Monoid a => a
mempty
            }

distributeSingleStm ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  DistAcc rep ->
  Stm SOACS ->
  DistNestT
    rep
    m
    ( Maybe
        ( PostStms rep,
          Result,
          KernelNest,
          DistAcc rep
        )
    )
distributeSingleStm :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep
-> Stm SOACS
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
distributeSingleStm DistAcc rep
acc Stm SOACS
bnd = do
  Nestings
nest <- (DistEnv rep m -> Nestings) -> DistNestT rep m Nestings
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Nestings
forall rep (m :: * -> *). DistEnv rep m -> Nestings
distNest
  [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl <- DistNestT
  rep
  m
  ([SubExp]
   -> [Char]
   -> ThreadRecommendation
   -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistNestT rep m (MkSegLevel rep (DistNestT rep m))
mkSegLevel
  ([SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
-> Nestings
-> Targets
-> Stms rep
-> DistNestT rep m (Maybe (Targets, Stms rep))
forall rep (m :: * -> *).
(DistRep rep, MonadFreshNames m, LocalScope rep m,
 MonadLogger m) =>
MkSegLevel rep m
-> Nestings -> Targets -> Stms rep -> m (Maybe (Targets, Stms rep))
tryDistribute [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl Nestings
nest (DistAcc rep -> Targets
forall rep. DistAcc rep -> Targets
distTargets DistAcc rep
acc) (DistAcc rep -> Stms rep
forall rep. DistAcc rep -> Stms rep
distStms DistAcc rep
acc) DistNestT rep m (Maybe (Targets, Stms rep))
-> (Maybe (Targets, Stms rep)
    -> DistNestT
         rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)))
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe (Targets, Stms rep)
Nothing -> Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
forall a. Maybe a
Nothing
    Just (Targets
targets, Stms rep
distributed_bnds) ->
      Nestings
-> Targets
-> Stm SOACS
-> DistNestT rep m (Maybe ([SubExp], Targets, KernelNest))
forall (m :: * -> *) t rep.
(MonadFreshNames m, HasScope t m, ASTRep rep) =>
Nestings
-> Targets -> Stm rep -> m (Maybe ([SubExp], Targets, KernelNest))
tryDistributeStm Nestings
nest Targets
targets Stm SOACS
bnd DistNestT rep m (Maybe ([SubExp], Targets, KernelNest))
-> (Maybe ([SubExp], Targets, KernelNest)
    -> DistNestT
         rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)))
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe ([SubExp], Targets, KernelNest)
Nothing -> Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
forall a. Maybe a
Nothing
        Just ([SubExp]
res, Targets
targets', KernelNest
new_kernel_nest) ->
          Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
 -> DistNestT
      rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)))
-> Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
-> DistNestT
     rep m (Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep))
forall a b. (a -> b) -> a -> b
$
            (PostStms rep, [SubExp], KernelNest, DistAcc rep)
-> Maybe (PostStms rep, [SubExp], KernelNest, DistAcc rep)
forall a. a -> Maybe a
Just
              ( Stms rep -> PostStms rep
forall rep. Stms rep -> PostStms rep
PostStms Stms rep
distributed_bnds,
                [SubExp]
res,
                KernelNest
new_kernel_nest,
                DistAcc :: forall rep. Targets -> Stms rep -> DistAcc rep
DistAcc
                  { distTargets :: Targets
distTargets = Targets
targets',
                    distStms :: Stms rep
distStms = Stms rep
forall a. Monoid a => a
mempty
                  }
              )

segmentedScatterKernel ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  KernelNest ->
  [Int] ->
  PatternT Type ->
  Certificates ->
  SubExp ->
  Lambda rep ->
  [VName] ->
  [(Shape, Int, VName)] ->
  DistNestT rep m (Stms rep)
segmentedScatterKernel :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> PatternT Type
-> Certificates
-> SubExp
-> Lambda rep
-> [VName]
-> [(ShapeBase SubExp, Int, VName)]
-> DistNestT rep m (Stms rep)
segmentedScatterKernel KernelNest
nest [Int]
perm PatternT Type
scatter_pat Certificates
cs SubExp
scatter_w Lambda rep
lam [VName]
ivs [(ShapeBase SubExp, Int, VName)]
dests = do
  -- We replicate some of the checking done by 'isSegmentedOp', but
  -- things are different because a scatter is not a reduction or
  -- scan.
  --
  -- First, pretend that the scatter is also part of the nesting.  The
  -- KernelNest we produce here is technically not sensible, but it's
  -- good enough for flatKernel to work.
  let nesting :: LoopNesting
nesting =
        PatternT Type
-> StmAux () -> SubExp -> [(Param Type, VName)] -> LoopNesting
MapNesting PatternT Type
scatter_pat (Certificates -> Attrs -> () -> StmAux ()
forall dec. Certificates -> Attrs -> dec -> StmAux dec
StmAux Certificates
cs Attrs
forall a. Monoid a => a
mempty ()) SubExp
scatter_w ([(Param Type, VName)] -> LoopNesting)
-> [(Param Type, VName)] -> LoopNesting
forall a b. (a -> b) -> a -> b
$ [Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda rep -> [LParam rep]
forall rep. LambdaT rep -> [LParam rep]
lambdaParams Lambda rep
lam) [VName]
ivs
      nest' :: KernelNest
nest' =
        (PatternT Type, [SubExp])
-> LoopNesting -> KernelNest -> KernelNest
pushInnerKernelNesting (PatternT Type
scatter_pat, BodyT rep -> [SubExp]
forall rep. BodyT rep -> [SubExp]
bodyResult (BodyT rep -> [SubExp]) -> BodyT rep -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> BodyT rep
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda rep
lam) LoopNesting
nesting KernelNest
nest
  ([(VName, SubExp)]
ispace, [KernelInput]
kernel_inps) <- KernelNest -> DistNestT rep m ([(VName, SubExp)], [KernelInput])
forall (m :: * -> *).
MonadFreshNames m =>
KernelNest -> m ([(VName, SubExp)], [KernelInput])
flatKernel KernelNest
nest'

  let ([ShapeBase SubExp]
as_ws, [Int]
as_ns, [VName]
as) = [(ShapeBase SubExp, Int, VName)]
-> ([ShapeBase SubExp], [Int], [VName])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(ShapeBase SubExp, Int, VName)]
dests
      indexes :: [Int]
indexes = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
as_ns ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (ShapeBase SubExp -> Int) -> [ShapeBase SubExp] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ShapeBase SubExp -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ShapeBase SubExp]
as_ws

  -- The input/output arrays ('as') _must_ correspond to some kernel
  -- input, or else the original nested scatter would have been
  -- ill-typed.  Find them.
  [KernelInput]
as_inps <- (VName -> DistNestT rep m KernelInput)
-> [VName] -> DistNestT rep m [KernelInput]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([KernelInput] -> VName -> DistNestT rep m KernelInput
forall {m :: * -> *} {t :: * -> *}.
(Monad m, Foldable t) =>
t KernelInput -> VName -> m KernelInput
findInput [KernelInput]
kernel_inps) [VName]
as

  [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl <- DistNestT
  rep
  m
  ([SubExp]
   -> [Char]
   -> ThreadRecommendation
   -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistNestT rep m (MkSegLevel rep (DistNestT rep m))
mkSegLevel

  let rts :: [Type]
rts =
        ([Type] -> [Type]) -> [[Type]] -> [Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
take Int
1) ([[Type]] -> [Type]) -> [[Type]] -> [Type]
forall a b. (a -> b) -> a -> b
$
          [Int] -> [Type] -> [[Type]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
as_ns ([Type] -> [[Type]]) -> [Type] -> [[Type]]
forall a b. (a -> b) -> a -> b
$
            Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
indexes) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda rep
lam
      ([SubExp]
is, [SubExp]
vs) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
indexes) ([SubExp] -> ([SubExp], [SubExp]))
-> [SubExp] -> ([SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ BodyT rep -> [SubExp]
forall rep. BodyT rep -> [SubExp]
bodyResult (BodyT rep -> [SubExp]) -> BodyT rep -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> BodyT rep
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda rep
lam

  -- Maybe add certificates to the indices.
  ([SubExp]
is', Stms rep
k_body_stms) <- Binder rep [SubExp] -> DistNestT rep m ([SubExp], Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Binder rep [SubExp] -> DistNestT rep m ([SubExp], Stms rep))
-> Binder rep [SubExp] -> DistNestT rep m ([SubExp], Stms rep)
forall a b. (a -> b) -> a -> b
$ do
    Stms (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms (Stms (Rep (BinderT rep (State VNameSource)))
 -> BinderT rep (State VNameSource) ())
-> Stms (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BodyT rep -> Stms rep
forall rep. BodyT rep -> Stms rep
bodyStms (BodyT rep -> Stms rep) -> BodyT rep -> Stms rep
forall a b. (a -> b) -> a -> b
$ Lambda rep -> BodyT rep
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda rep
lam
    [SubExp]
-> (SubExp -> BinderT rep (State VNameSource) SubExp)
-> Binder rep [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [SubExp]
is ((SubExp -> BinderT rep (State VNameSource) SubExp)
 -> Binder rep [SubExp])
-> (SubExp -> BinderT rep (State VNameSource) SubExp)
-> Binder rep [SubExp]
forall a b. (a -> b) -> a -> b
$ \SubExp
i ->
      if Certificates
cs Certificates -> Certificates -> Bool
forall a. Eq a => a -> a -> Bool
== Certificates
forall a. Monoid a => a
mempty
        then SubExp -> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
i
        else Certificates
-> BinderT rep (State VNameSource) SubExp
-> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (BinderT rep (State VNameSource) SubExp
 -> BinderT rep (State VNameSource) SubExp)
-> BinderT rep (State VNameSource) SubExp
-> BinderT rep (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ [Char]
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"scatter_i" (Exp (Rep (BinderT rep (State VNameSource)))
 -> BinderT rep (State VNameSource) SubExp)
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT rep) -> BasicOp -> ExpT rep
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
i

  let k_body :: KernelBody rep
k_body =
        [(ShapeBase SubExp, Int, KernelInput)]
-> [SubExp]
-> [(ShapeBase SubExp, KernelInput, [([SubExp], SubExp)])]
forall array a.
[(ShapeBase SubExp, Int, array)]
-> [a] -> [(ShapeBase SubExp, array, [([a], a)])]
groupScatterResults ([ShapeBase SubExp]
-> [Int] -> [KernelInput] -> [(ShapeBase SubExp, Int, KernelInput)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [ShapeBase SubExp]
as_ws [Int]
as_ns [KernelInput]
as_inps) ([SubExp]
is' [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
vs)
          [(ShapeBase SubExp, KernelInput, [([SubExp], SubExp)])]
-> ([(ShapeBase SubExp, KernelInput, [([SubExp], SubExp)])]
    -> [KernelResult])
-> [KernelResult]
forall a b. a -> (a -> b) -> b
& ((ShapeBase SubExp, KernelInput, [([SubExp], SubExp)])
 -> KernelResult)
-> [(ShapeBase SubExp, KernelInput, [([SubExp], SubExp)])]
-> [KernelResult]
forall a b. (a -> b) -> [a] -> [b]
map ([(VName, SubExp)]
-> (ShapeBase SubExp, KernelInput, [([SubExp], SubExp)])
-> KernelResult
inPlaceReturn [(VName, SubExp)]
ispace)
          [KernelResult]
-> ([KernelResult] -> KernelBody rep) -> KernelBody rep
forall a b. a -> (a -> b) -> b
& BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
forall rep.
BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
KernelBody () Stms rep
k_body_stms
      -- Remove unused kernel inputs, since some of these might
      -- reference the array we are scattering into.
      kernel_inps' :: [KernelInput]
kernel_inps' =
        (KernelInput -> Bool) -> [KernelInput] -> [KernelInput]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Names -> Bool
`nameIn` KernelBody rep -> Names
forall a. FreeIn a => a -> Names
freeIn KernelBody rep
k_body) (VName -> Bool) -> (KernelInput -> VName) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelInput -> VName
kernelInputName) [KernelInput]
kernel_inps

  (SegOp (SegOpLevel rep) rep
k, Stms rep
k_bnds) <- ([SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
-> [(VName, SubExp)]
-> [KernelInput]
-> [Type]
-> KernelBody rep
-> DistNestT rep m (SegOp (SegOpLevel rep) rep, Stms rep)
forall rep (m :: * -> *).
(DistRep rep, HasScope rep m, MonadFreshNames m) =>
MkSegLevel rep m
-> [(VName, SubExp)]
-> [KernelInput]
-> [Type]
-> KernelBody rep
-> m (SegOp (SegOpLevel rep) rep, Stms rep)
mapKernel [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl [(VName, SubExp)]
ispace [KernelInput]
kernel_inps' [Type]
rts KernelBody rep
k_body

  (Stm rep -> DistNestT rep m (Stm rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Stm rep -> DistNestT rep m (Stm rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Stm rep -> m (Stm rep)
renameStm (Stms rep -> DistNestT rep m (Stms rep))
-> (BinderT rep (State VNameSource) ()
    -> DistNestT rep m (Stms rep))
-> BinderT rep (State VNameSource) ()
-> DistNestT rep m (Stms rep)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< BinderT rep (State VNameSource) () -> DistNestT rep m (Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (Stms rep)
runBinder_ (BinderT rep (State VNameSource) () -> DistNestT rep m (Stms rep))
-> BinderT rep (State VNameSource) () -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$ do
    Stms (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms rep
Stms (Rep (BinderT rep (State VNameSource)))
k_bnds

    let pat :: PatternT Type
pat =
          [PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] ([PatElemT Type] -> PatternT Type)
-> [PatElemT Type] -> PatternT Type
forall a b. (a -> b) -> a -> b
$
            [Int] -> [PatElemT Type] -> [PatElemT Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([PatElemT Type] -> [PatElemT Type])
-> [PatElemT Type] -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$
              PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements (PatternT Type -> [PatElemT Type])
-> PatternT Type -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$ LoopNesting -> PatternT Type
loopNestingPattern (LoopNesting -> PatternT Type) -> LoopNesting -> PatternT Type
forall a b. (a -> b) -> a -> b
$ KernelNest -> LoopNesting
forall a b. (a, b) -> a
fst KernelNest
nest

    Pattern (Rep (BinderT rep (State VNameSource)))
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Rep m) -> Exp (Rep m) -> m ()
letBind PatternT Type
Pattern (Rep (BinderT rep (State VNameSource)))
pat (Exp (Rep (BinderT rep (State VNameSource)))
 -> BinderT rep (State VNameSource) ())
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ Op rep -> ExpT rep
forall rep. Op rep -> ExpT rep
Op (Op rep -> ExpT rep) -> Op rep -> ExpT rep
forall a b. (a -> b) -> a -> b
$ SegOp (SegOpLevel rep) rep -> Op rep
forall rep. HasSegOp rep => SegOp (SegOpLevel rep) rep -> Op rep
segOp SegOp (SegOpLevel rep) rep
k
  where
    findInput :: t KernelInput -> VName -> m KernelInput
findInput t KernelInput
kernel_inps VName
a =
      m KernelInput
-> (KernelInput -> m KernelInput)
-> Maybe KernelInput
-> m KernelInput
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m KernelInput
forall {a}. a
bad KernelInput -> m KernelInput
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe KernelInput -> m KernelInput)
-> Maybe KernelInput -> m KernelInput
forall a b. (a -> b) -> a -> b
$ (KernelInput -> Bool) -> t KernelInput -> Maybe KernelInput
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
a) (VName -> Bool) -> (KernelInput -> VName) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelInput -> VName
kernelInputName) t KernelInput
kernel_inps
    bad :: a
bad = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Ill-typed nested scatter encountered."

    inPlaceReturn :: [(VName, SubExp)]
-> (ShapeBase SubExp, KernelInput, [([SubExp], SubExp)])
-> KernelResult
inPlaceReturn [(VName, SubExp)]
ispace (ShapeBase SubExp
aw, KernelInput
inp, [([SubExp], SubExp)]
is_vs) =
      ShapeBase SubExp
-> VName -> [(Slice SubExp, SubExp)] -> KernelResult
WriteReturns
        ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> [SubExp]
forall a. [a] -> [a]
init [SubExp]
ws [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
aw))
        (KernelInput -> VName
kernelInputArray KernelInput
inp)
        [((SubExp -> DimIndex SubExp) -> [SubExp] -> Slice SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix ([SubExp] -> Slice SubExp) -> [SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var ([VName] -> [VName]
forall a. [a] -> [a]
init [VName]
gtids) [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
is, SubExp
v) | ([SubExp]
is, SubExp
v) <- [([SubExp], SubExp)]
is_vs]
      where
        ([VName]
gtids, [SubExp]
ws) = [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(VName, SubExp)]
ispace

segmentedUpdateKernel ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  KernelNest ->
  [Int] ->
  Certificates ->
  VName ->
  Slice SubExp ->
  VName ->
  DistNestT rep m (Stms rep)
segmentedUpdateKernel :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Certificates
-> VName
-> Slice SubExp
-> VName
-> DistNestT rep m (Stms rep)
segmentedUpdateKernel KernelNest
nest [Int]
perm Certificates
cs VName
arr Slice SubExp
slice VName
v = do
  ([(VName, SubExp)]
base_ispace, [KernelInput]
kernel_inps) <- KernelNest -> DistNestT rep m ([(VName, SubExp)], [KernelInput])
forall (m :: * -> *).
MonadFreshNames m =>
KernelNest -> m ([(VName, SubExp)], [KernelInput])
flatKernel KernelNest
nest
  let slice_dims :: [SubExp]
slice_dims = Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
  [VName]
slice_gtids <- Int -> DistNestT rep m VName -> DistNestT rep m [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
slice_dims) ([Char] -> DistNestT rep m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"gtid_slice")

  let ispace :: [(VName, SubExp)]
ispace = [(VName, SubExp)]
base_ispace [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
forall a. [a] -> [a] -> [a]
++ [VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
slice_gtids [SubExp]
slice_dims

  ((Type
res_t, KernelResult
res), Stms rep
kstms) <- Binder rep (Type, KernelResult)
-> DistNestT rep m ((Type, KernelResult), Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Binder rep (Type, KernelResult)
 -> DistNestT rep m ((Type, KernelResult), Stms rep))
-> Binder rep (Type, KernelResult)
-> DistNestT rep m ((Type, KernelResult), Stms rep)
forall a b. (a -> b) -> a -> b
$ do
    -- Compute indexes into full array.
    SubExp
v' <-
      Certificates
-> BinderT rep (State VNameSource) SubExp
-> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (BinderT rep (State VNameSource) SubExp
 -> BinderT rep (State VNameSource) SubExp)
-> BinderT rep (State VNameSource) SubExp
-> BinderT rep (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$
        [Char]
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"v" (Exp (Rep (BinderT rep (State VNameSource)))
 -> BinderT rep (State VNameSource) SubExp)
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT rep) -> BasicOp -> ExpT rep
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
v (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ (VName -> DimIndex SubExp) -> [VName] -> Slice SubExp
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp)
-> (VName -> SubExp) -> VName -> DimIndex SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) [VName]
slice_gtids
    [SubExp]
slice_is <-
      (TPrimExp Int64 VName -> BinderT rep (State VNameSource) SubExp)
-> [TPrimExp Int64 VName]
-> BinderT rep (State VNameSource) [SubExp]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ([Char]
-> TPrimExp Int64 VName -> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *) a.
(MonadBinder m, ToExp a) =>
[Char] -> a -> m SubExp
toSubExp [Char]
"index") ([TPrimExp Int64 VName]
 -> BinderT rep (State VNameSource) [SubExp])
-> [TPrimExp Int64 VName]
-> BinderT rep (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$
        Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice ((DimIndex SubExp -> DimIndex (TPrimExp Int64 VName))
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> TPrimExp Int64 VName)
-> DimIndex SubExp -> DimIndex (TPrimExp Int64 VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) Slice SubExp
slice) ([TPrimExp Int64 VName] -> [TPrimExp Int64 VName])
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ (VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> TPrimExp Int64 VName
pe64 (SubExp -> TPrimExp Int64 VName)
-> (VName -> SubExp) -> VName -> TPrimExp Int64 VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) [VName]
slice_gtids

    let write_is :: [SubExp]
write_is = ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> ((VName, SubExp) -> VName) -> (VName, SubExp) -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst) [(VName, SubExp)]
base_ispace [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
slice_is
        arr' :: VName
arr' =
          VName -> (KernelInput -> VName) -> Maybe KernelInput -> VName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> VName
forall a. HasCallStack => [Char] -> a
error [Char]
"incorrectly typed Update") KernelInput -> VName
kernelInputArray (Maybe KernelInput -> VName) -> Maybe KernelInput -> VName
forall a b. (a -> b) -> a -> b
$
            (KernelInput -> Bool) -> [KernelInput] -> Maybe KernelInput
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
arr) (VName -> Bool) -> (KernelInput -> VName) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelInput -> VName
kernelInputName) [KernelInput]
kernel_inps
    Type
arr_t <- VName -> BinderT rep (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr'
    Type
v_t <- SubExp -> BinderT rep (State VNameSource) Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
v'
    (Type, KernelResult) -> Binder rep (Type, KernelResult)
forall (m :: * -> *) a. Monad m => a -> m a
return
      ( Type
v_t,
        ShapeBase SubExp
-> VName -> [(Slice SubExp, SubExp)] -> KernelResult
WriteReturns (Type -> ShapeBase SubExp
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
arr_t) VName
arr' [((SubExp -> DimIndex SubExp) -> [SubExp] -> Slice SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix [SubExp]
write_is, SubExp
v')]
      )

  -- Remove unused kernel inputs, since some of these might
  -- reference the array we are scattering into.
  let kernel_inps' :: [KernelInput]
kernel_inps' =
        (KernelInput -> Bool) -> [KernelInput] -> [KernelInput]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Names -> Bool
`nameIn` Stms rep -> Names
forall a. FreeIn a => a -> Names
freeIn Stms rep
kstms) (VName -> Bool) -> (KernelInput -> VName) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelInput -> VName
kernelInputName) [KernelInput]
kernel_inps

  [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl <- DistNestT
  rep
  m
  ([SubExp]
   -> [Char]
   -> ThreadRecommendation
   -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistNestT rep m (MkSegLevel rep (DistNestT rep m))
mkSegLevel
  (SegOp (SegOpLevel rep) rep
k, Stms rep
prestms) <-
    ([SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
-> [(VName, SubExp)]
-> [KernelInput]
-> [Type]
-> KernelBody rep
-> DistNestT rep m (SegOp (SegOpLevel rep) rep, Stms rep)
forall rep (m :: * -> *).
(DistRep rep, HasScope rep m, MonadFreshNames m) =>
MkSegLevel rep m
-> [(VName, SubExp)]
-> [KernelInput]
-> [Type]
-> KernelBody rep
-> m (SegOp (SegOpLevel rep) rep, Stms rep)
mapKernel [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl [(VName, SubExp)]
ispace [KernelInput]
kernel_inps' [Type
res_t] (KernelBody rep
 -> DistNestT rep m (SegOp (SegOpLevel rep) rep, Stms rep))
-> KernelBody rep
-> DistNestT rep m (SegOp (SegOpLevel rep) rep, Stms rep)
forall a b. (a -> b) -> a -> b
$
      BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
forall rep.
BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
KernelBody () Stms rep
kstms [KernelResult
res]

  (Stm rep -> DistNestT rep m (Stm rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Stm rep -> DistNestT rep m (Stm rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Stm rep -> m (Stm rep)
renameStm (Stms rep -> DistNestT rep m (Stms rep))
-> (Binder rep () -> DistNestT rep m (Stms rep))
-> Binder rep ()
-> DistNestT rep m (Stms rep)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Binder rep () -> DistNestT rep m (Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (Stms rep)
runBinder_ (Binder rep () -> DistNestT rep m (Stms rep))
-> Binder rep () -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$ do
    Stms (Rep (BinderT rep (State VNameSource))) -> Binder rep ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms rep
Stms (Rep (BinderT rep (State VNameSource)))
prestms

    let pat :: PatternT Type
pat =
          [PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] ([PatElemT Type] -> PatternT Type)
-> [PatElemT Type] -> PatternT Type
forall a b. (a -> b) -> a -> b
$
            [Int] -> [PatElemT Type] -> [PatElemT Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([PatElemT Type] -> [PatElemT Type])
-> [PatElemT Type] -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$
              PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements (PatternT Type -> [PatElemT Type])
-> PatternT Type -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$ LoopNesting -> PatternT Type
loopNestingPattern (LoopNesting -> PatternT Type) -> LoopNesting -> PatternT Type
forall a b. (a -> b) -> a -> b
$ KernelNest -> LoopNesting
forall a b. (a, b) -> a
fst KernelNest
nest

    Pattern (Rep (BinderT rep (State VNameSource)))
-> Exp (Rep (BinderT rep (State VNameSource))) -> Binder rep ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Rep m) -> Exp (Rep m) -> m ()
letBind PatternT Type
Pattern (Rep (BinderT rep (State VNameSource)))
pat (Exp (Rep (BinderT rep (State VNameSource))) -> Binder rep ())
-> Exp (Rep (BinderT rep (State VNameSource))) -> Binder rep ()
forall a b. (a -> b) -> a -> b
$ Op rep -> ExpT rep
forall rep. Op rep -> ExpT rep
Op (Op rep -> ExpT rep) -> Op rep -> ExpT rep
forall a b. (a -> b) -> a -> b
$ SegOp (SegOpLevel rep) rep -> Op rep
forall rep. HasSegOp rep => SegOp (SegOpLevel rep) rep -> Op rep
segOp SegOp (SegOpLevel rep) rep
k

segmentedGatherKernel ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  KernelNest ->
  Certificates ->
  VName ->
  Slice SubExp ->
  DistNestT rep m (Stms rep)
segmentedGatherKernel :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> Certificates
-> VName
-> Slice SubExp
-> DistNestT rep m (Stms rep)
segmentedGatherKernel KernelNest
nest Certificates
cs VName
arr Slice SubExp
slice = do
  let slice_dims :: [SubExp]
slice_dims = Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
  [VName]
slice_gtids <- Int -> DistNestT rep m VName -> DistNestT rep m [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
slice_dims) ([Char] -> DistNestT rep m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"gtid_slice")

  ([(VName, SubExp)]
base_ispace, [KernelInput]
kernel_inps) <- KernelNest -> DistNestT rep m ([(VName, SubExp)], [KernelInput])
forall (m :: * -> *).
MonadFreshNames m =>
KernelNest -> m ([(VName, SubExp)], [KernelInput])
flatKernel KernelNest
nest
  let ispace :: [(VName, SubExp)]
ispace = [(VName, SubExp)]
base_ispace [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
forall a. [a] -> [a] -> [a]
++ [VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
slice_gtids [SubExp]
slice_dims

  ((Type
res_t, KernelResult
res), Stms rep
kstms) <- Binder rep (Type, KernelResult)
-> DistNestT rep m ((Type, KernelResult), Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Binder rep (Type, KernelResult)
 -> DistNestT rep m ((Type, KernelResult), Stms rep))
-> Binder rep (Type, KernelResult)
-> DistNestT rep m ((Type, KernelResult), Stms rep)
forall a b. (a -> b) -> a -> b
$ do
    -- Compute indexes into full array.
    Slice SubExp
slice'' <-
      Slice (TPrimExp Int64 VName)
-> BinderT rep (State VNameSource) (Slice SubExp)
forall (m :: * -> *).
MonadBinder m =>
Slice (TPrimExp Int64 VName) -> m (Slice SubExp)
subExpSlice (Slice (TPrimExp Int64 VName)
 -> BinderT rep (State VNameSource) (Slice SubExp))
-> Slice (TPrimExp Int64 VName)
-> BinderT rep (State VNameSource) (Slice SubExp)
forall a b. (a -> b) -> a -> b
$
        Slice (TPrimExp Int64 VName)
-> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
forall d. Num d => Slice d -> Slice d -> Slice d
sliceSlice (Slice SubExp -> Slice (TPrimExp Int64 VName)
primExpSlice Slice SubExp
slice) (Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName))
-> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$
          Slice SubExp -> Slice (TPrimExp Int64 VName)
primExpSlice (Slice SubExp -> Slice (TPrimExp Int64 VName))
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ (VName -> DimIndex SubExp) -> [VName] -> Slice SubExp
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp)
-> (VName -> SubExp) -> VName -> DimIndex SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) [VName]
slice_gtids
    SubExp
v' <- Certificates
-> BinderT rep (State VNameSource) SubExp
-> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (BinderT rep (State VNameSource) SubExp
 -> BinderT rep (State VNameSource) SubExp)
-> BinderT rep (State VNameSource) SubExp
-> BinderT rep (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ [Char]
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"v" (Exp (Rep (BinderT rep (State VNameSource)))
 -> BinderT rep (State VNameSource) SubExp)
-> Exp (Rep (BinderT rep (State VNameSource)))
-> BinderT rep (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT rep) -> BasicOp -> ExpT rep
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
arr Slice SubExp
slice''
    Type
v_t <- SubExp -> BinderT rep (State VNameSource) Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
v'
    (Type, KernelResult) -> Binder rep (Type, KernelResult)
forall (m :: * -> *) a. Monad m => a -> m a
return (Type
v_t, ResultManifest -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify SubExp
v')

  [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl <- DistNestT
  rep
  m
  ([SubExp]
   -> [Char]
   -> ThreadRecommendation
   -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistNestT rep m (MkSegLevel rep (DistNestT rep m))
mkSegLevel
  (SegOp (SegOpLevel rep) rep
k, Stms rep
prestms) <-
    ([SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep (DistNestT rep m) (SegOpLevel rep))
-> [(VName, SubExp)]
-> [KernelInput]
-> [Type]
-> KernelBody rep
-> DistNestT rep m (SegOp (SegOpLevel rep) rep, Stms rep)
forall rep (m :: * -> *).
(DistRep rep, HasScope rep m, MonadFreshNames m) =>
MkSegLevel rep m
-> [(VName, SubExp)]
-> [KernelInput]
-> [Type]
-> KernelBody rep
-> m (SegOp (SegOpLevel rep) rep, Stms rep)
mapKernel [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep (DistNestT rep m) (SegOpLevel rep)
mk_lvl [(VName, SubExp)]
ispace [KernelInput]
kernel_inps [Type
res_t] (KernelBody rep
 -> DistNestT rep m (SegOp (SegOpLevel rep) rep, Stms rep))
-> KernelBody rep
-> DistNestT rep m (SegOp (SegOpLevel rep) rep, Stms rep)
forall a b. (a -> b) -> a -> b
$
      BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
forall rep.
BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
KernelBody () Stms rep
kstms [KernelResult
res]

  (Stm rep -> DistNestT rep m (Stm rep))
-> Stms rep -> DistNestT rep m (Stms rep)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Stm rep -> DistNestT rep m (Stm rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Stm rep -> m (Stm rep)
renameStm (Stms rep -> DistNestT rep m (Stms rep))
-> (Binder rep () -> DistNestT rep m (Stms rep))
-> Binder rep ()
-> DistNestT rep m (Stms rep)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Binder rep () -> DistNestT rep m (Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (Stms rep)
runBinder_ (Binder rep () -> DistNestT rep m (Stms rep))
-> Binder rep () -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$ do
    Stms (Rep (BinderT rep (State VNameSource))) -> Binder rep ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms rep
Stms (Rep (BinderT rep (State VNameSource)))
prestms

    let pat :: PatternT Type
pat = [PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] ([PatElemT Type] -> PatternT Type)
-> [PatElemT Type] -> PatternT Type
forall a b. (a -> b) -> a -> b
$ PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements (PatternT Type -> [PatElemT Type])
-> PatternT Type -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$ LoopNesting -> PatternT Type
loopNestingPattern (LoopNesting -> PatternT Type) -> LoopNesting -> PatternT Type
forall a b. (a -> b) -> a -> b
$ KernelNest -> LoopNesting
forall a b. (a, b) -> a
fst KernelNest
nest

    Pattern (Rep (BinderT rep (State VNameSource)))
-> Exp (Rep (BinderT rep (State VNameSource))) -> Binder rep ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Rep m) -> Exp (Rep m) -> m ()
letBind PatternT Type
Pattern (Rep (BinderT rep (State VNameSource)))
pat (Exp (Rep (BinderT rep (State VNameSource))) -> Binder rep ())
-> Exp (Rep (BinderT rep (State VNameSource))) -> Binder rep ()
forall a b. (a -> b) -> a -> b
$ Op rep -> ExpT rep
forall rep. Op rep -> ExpT rep
Op (Op rep -> ExpT rep) -> Op rep -> ExpT rep
forall a b. (a -> b) -> a -> b
$ SegOp (SegOpLevel rep) rep -> Op rep
forall rep. HasSegOp rep => SegOp (SegOpLevel rep) rep -> Op rep
segOp SegOp (SegOpLevel rep) rep
k

segmentedHistKernel ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  KernelNest ->
  [Int] ->
  Certificates ->
  SubExp ->
  [SOACS.HistOp SOACS] ->
  Lambda rep ->
  [VName] ->
  DistNestT rep m (Stms rep)
segmentedHistKernel :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Certificates
-> SubExp
-> [HistOp SOACS]
-> Lambda rep
-> [VName]
-> DistNestT rep m (Stms rep)
segmentedHistKernel KernelNest
nest [Int]
perm Certificates
cs SubExp
hist_w [HistOp SOACS]
ops Lambda rep
lam [VName]
arrs = do
  -- We replicate some of the checking done by 'isSegmentedOp', but
  -- things are different because a Hist is not a reduction or
  -- scan.
  ([(VName, SubExp)]
ispace, [KernelInput]
inputs) <- KernelNest -> DistNestT rep m ([(VName, SubExp)], [KernelInput])
forall (m :: * -> *).
MonadFreshNames m =>
KernelNest -> m ([(VName, SubExp)], [KernelInput])
flatKernel KernelNest
nest
  let orig_pat :: PatternT Type
orig_pat =
        [PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] ([PatElemT Type] -> PatternT Type)
-> [PatElemT Type] -> PatternT Type
forall a b. (a -> b) -> a -> b
$
          [Int] -> [PatElemT Type] -> [PatElemT Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([PatElemT Type] -> [PatElemT Type])
-> [PatElemT Type] -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$
            PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements (PatternT Type -> [PatElemT Type])
-> PatternT Type -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$ LoopNesting -> PatternT Type
loopNestingPattern (LoopNesting -> PatternT Type) -> LoopNesting -> PatternT Type
forall a b. (a -> b) -> a -> b
$ KernelNest -> LoopNesting
forall a b. (a, b) -> a
fst KernelNest
nest

  -- The input/output arrays _must_ correspond to some kernel input,
  -- or else the original nested Hist would have been ill-typed.
  -- Find them.
  [HistOp SOACS]
ops' <- [HistOp SOACS]
-> (HistOp SOACS -> DistNestT rep m (HistOp SOACS))
-> DistNestT rep m [HistOp SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [HistOp SOACS]
ops ((HistOp SOACS -> DistNestT rep m (HistOp SOACS))
 -> DistNestT rep m [HistOp SOACS])
-> (HistOp SOACS -> DistNestT rep m (HistOp SOACS))
-> DistNestT rep m [HistOp SOACS]
forall a b. (a -> b) -> a -> b
$ \(SOACS.HistOp SubExp
num_bins SubExp
rf [VName]
dests [SubExp]
nes Lambda
op) ->
    SubExp -> SubExp -> [VName] -> [SubExp] -> Lambda -> HistOp SOACS
forall rep.
SubExp -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
SOACS.HistOp SubExp
num_bins SubExp
rf
      ([VName] -> [SubExp] -> Lambda -> HistOp SOACS)
-> DistNestT rep m [VName]
-> DistNestT rep m ([SubExp] -> Lambda -> HistOp SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> DistNestT rep m VName)
-> [VName] -> DistNestT rep m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((KernelInput -> VName)
-> DistNestT rep m KernelInput -> DistNestT rep m VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap KernelInput -> VName
kernelInputArray (DistNestT rep m KernelInput -> DistNestT rep m VName)
-> (VName -> DistNestT rep m KernelInput)
-> VName
-> DistNestT rep m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [KernelInput] -> VName -> DistNestT rep m KernelInput
forall {m :: * -> *} {t :: * -> *}.
(Monad m, Foldable t) =>
t KernelInput -> VName -> m KernelInput
findInput [KernelInput]
inputs) [VName]
dests
      DistNestT rep m ([SubExp] -> Lambda -> HistOp SOACS)
-> DistNestT rep m [SubExp]
-> DistNestT rep m (Lambda -> HistOp SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [SubExp] -> DistNestT rep m [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
nes
      DistNestT rep m (Lambda -> HistOp SOACS)
-> DistNestT rep m Lambda -> DistNestT rep m (HistOp SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Lambda -> DistNestT rep m Lambda
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda
op

  [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl <- (DistEnv rep m
 -> [SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep m (SegOpLevel rep))
-> DistNestT
     rep
     m
     ([SubExp]
      -> [Char]
      -> ThreadRecommendation
      -> BinderT rep m (SegOpLevel rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m
-> [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep m (SegOpLevel rep)
forall rep (m :: * -> *). DistEnv rep m -> MkSegLevel rep m
distSegLevel
  Lambda -> Binder rep (Lambda rep)
onLambda <- (DistEnv rep m -> Lambda -> Binder rep (Lambda rep))
-> DistNestT rep m (Lambda -> Binder rep (Lambda rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m -> Lambda -> Binder rep (Lambda rep)
forall rep (m :: * -> *).
DistEnv rep m -> Lambda -> Binder rep (Lambda rep)
distOnSOACSLambda
  let onLambda' :: Lambda -> BinderT rep m (Lambda rep)
onLambda' = ((Lambda rep, Stms rep) -> Lambda rep)
-> BinderT rep m (Lambda rep, Stms rep)
-> BinderT rep m (Lambda rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Lambda rep, Stms rep) -> Lambda rep
forall a b. (a, b) -> a
fst (BinderT rep m (Lambda rep, Stms rep)
 -> BinderT rep m (Lambda rep))
-> (Lambda -> BinderT rep m (Lambda rep, Stms rep))
-> Lambda
-> BinderT rep m (Lambda rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binder rep (Lambda rep) -> BinderT rep m (Lambda rep, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Binder rep (Lambda rep) -> BinderT rep m (Lambda rep, Stms rep))
-> (Lambda -> Binder rep (Lambda rep))
-> Lambda
-> BinderT rep m (Lambda rep, Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda -> Binder rep (Lambda rep)
onLambda
  m (Stms rep) -> DistNestT rep m (Stms rep)
forall rep (m :: * -> *) a.
(LocalScope rep m, DistRep rep) =>
m a -> DistNestT rep m a
liftInner (m (Stms rep) -> DistNestT rep m (Stms rep))
-> m (Stms rep) -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$
    BinderT rep m () -> m (Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
BinderT rep m a -> m (Stms rep)
runBinderT'_ (BinderT rep m () -> m (Stms rep))
-> BinderT rep m () -> m (Stms rep)
forall a b. (a -> b) -> a -> b
$ do
      -- It is important not to launch unnecessarily many threads for
      -- histograms, because it may mean we unnecessarily need to reduce
      -- subhistograms as well.
      SegOpLevel rep
lvl <- [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl (SubExp
hist_w SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
ispace) [Char]
"seghist" (ThreadRecommendation -> BinderT rep m (SegOpLevel rep))
-> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
forall a b. (a -> b) -> a -> b
$ SegVirt -> ThreadRecommendation
NoRecommendation SegVirt
SegNoVirt
      Stms rep -> BinderT rep m ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms
        (Stms rep -> BinderT rep m ())
-> BinderT rep m (Stms rep) -> BinderT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Lambda -> BinderT rep m (Lambda (Rep (BinderT rep m))))
-> SegOpLevel (Rep (BinderT rep m))
-> PatternT Type
-> [(VName, SubExp)]
-> [KernelInput]
-> Certificates
-> SubExp
-> [HistOp SOACS]
-> Lambda (Rep (BinderT rep m))
-> [VName]
-> BinderT rep m (Stms (Rep (BinderT rep m)))
forall (m :: * -> *).
(MonadBinder m, DistRep (Rep m)) =>
(Lambda -> m (Lambda (Rep m)))
-> SegOpLevel (Rep m)
-> PatternT Type
-> [(VName, SubExp)]
-> [KernelInput]
-> Certificates
-> SubExp
-> [HistOp SOACS]
-> Lambda (Rep m)
-> [VName]
-> m (Stms (Rep m))
histKernel Lambda -> BinderT rep m (Lambda rep)
Lambda -> BinderT rep m (Lambda (Rep (BinderT rep m)))
onLambda' SegOpLevel rep
SegOpLevel (Rep (BinderT rep m))
lvl PatternT Type
orig_pat [(VName, SubExp)]
ispace [KernelInput]
inputs Certificates
cs SubExp
hist_w [HistOp SOACS]
ops' Lambda rep
Lambda (Rep (BinderT rep m))
lam [VName]
arrs
  where
    findInput :: t KernelInput -> VName -> m KernelInput
findInput t KernelInput
kernel_inps VName
a =
      m KernelInput
-> (KernelInput -> m KernelInput)
-> Maybe KernelInput
-> m KernelInput
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m KernelInput
forall {a}. a
bad KernelInput -> m KernelInput
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe KernelInput -> m KernelInput)
-> Maybe KernelInput -> m KernelInput
forall a b. (a -> b) -> a -> b
$ (KernelInput -> Bool) -> t KernelInput -> Maybe KernelInput
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
a) (VName -> Bool) -> (KernelInput -> VName) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelInput -> VName
kernelInputName) t KernelInput
kernel_inps
    bad :: a
bad = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Ill-typed nested Hist encountered."

histKernel ::
  (MonadBinder m, DistRep (Rep m)) =>
  (Lambda SOACS -> m (Lambda (Rep m))) ->
  SegOpLevel (Rep m) ->
  PatternT Type ->
  [(VName, SubExp)] ->
  [KernelInput] ->
  Certificates ->
  SubExp ->
  [SOACS.HistOp SOACS] ->
  Lambda (Rep m) ->
  [VName] ->
  m (Stms (Rep m))
histKernel :: forall (m :: * -> *).
(MonadBinder m, DistRep (Rep m)) =>
(Lambda -> m (Lambda (Rep m)))
-> SegOpLevel (Rep m)
-> PatternT Type
-> [(VName, SubExp)]
-> [KernelInput]
-> Certificates
-> SubExp
-> [HistOp SOACS]
-> Lambda (Rep m)
-> [VName]
-> m (Stms (Rep m))
histKernel Lambda -> m (Lambda (Rep m))
onLambda SegOpLevel (Rep m)
lvl PatternT Type
orig_pat [(VName, SubExp)]
ispace [KernelInput]
inputs Certificates
cs SubExp
hist_w [HistOp SOACS]
ops Lambda (Rep m)
lam [VName]
arrs = BinderT (Rep m) m () -> m (Seq (Stm (Rep m)))
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
BinderT rep m a -> m (Stms rep)
runBinderT'_ (BinderT (Rep m) m () -> m (Seq (Stm (Rep m))))
-> BinderT (Rep m) m () -> m (Seq (Stm (Rep m)))
forall a b. (a -> b) -> a -> b
$ do
  [HistOp (Rep m)]
ops' <- [HistOp SOACS]
-> (HistOp SOACS -> BinderT (Rep m) m (HistOp (Rep m)))
-> BinderT (Rep m) m [HistOp (Rep m)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [HistOp SOACS]
ops ((HistOp SOACS -> BinderT (Rep m) m (HistOp (Rep m)))
 -> BinderT (Rep m) m [HistOp (Rep m)])
-> (HistOp SOACS -> BinderT (Rep m) m (HistOp (Rep m)))
-> BinderT (Rep m) m [HistOp (Rep m)]
forall a b. (a -> b) -> a -> b
$ \(SOACS.HistOp SubExp
num_bins SubExp
rf [VName]
dests [SubExp]
nes Lambda
op) -> do
    (Lambda
op', [SubExp]
nes', ShapeBase SubExp
shape) <- Lambda
-> [SubExp]
-> BinderT (Rep m) m (Lambda, [SubExp], ShapeBase SubExp)
forall (m :: * -> *).
MonadBinder m =>
Lambda -> [SubExp] -> m (Lambda, [SubExp], ShapeBase SubExp)
determineReduceOp Lambda
op [SubExp]
nes
    Lambda (Rep m)
op'' <- m (Lambda (Rep m)) -> BinderT (Rep m) m (Lambda (Rep m))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (Lambda (Rep m)) -> BinderT (Rep m) m (Lambda (Rep m)))
-> m (Lambda (Rep m)) -> BinderT (Rep m) m (Lambda (Rep m))
forall a b. (a -> b) -> a -> b
$ Lambda -> m (Lambda (Rep m))
onLambda Lambda
op'
    HistOp (Rep m) -> BinderT (Rep m) m (HistOp (Rep m))
forall (m :: * -> *) a. Monad m => a -> m a
return (HistOp (Rep m) -> BinderT (Rep m) m (HistOp (Rep m)))
-> HistOp (Rep m) -> BinderT (Rep m) m (HistOp (Rep m))
forall a b. (a -> b) -> a -> b
$ SubExp
-> SubExp
-> [VName]
-> [SubExp]
-> ShapeBase SubExp
-> Lambda (Rep m)
-> HistOp (Rep m)
forall rep.
SubExp
-> SubExp
-> [VName]
-> [SubExp]
-> ShapeBase SubExp
-> Lambda rep
-> HistOp rep
HistOp SubExp
num_bins SubExp
rf [VName]
dests [SubExp]
nes' ShapeBase SubExp
shape Lambda (Rep m)
op''

  let isDest :: VName -> Bool
isDest = (VName -> [VName] -> Bool) -> [VName] -> VName -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem ([VName] -> VName -> Bool) -> [VName] -> VName -> Bool
forall a b. (a -> b) -> a -> b
$ (HistOp (Rep m) -> [VName]) -> [HistOp (Rep m)] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap HistOp (Rep m) -> [VName]
forall rep. HistOp rep -> [VName]
histDest [HistOp (Rep m)]
ops'
      inputs' :: [KernelInput]
inputs' = (KernelInput -> Bool) -> [KernelInput] -> [KernelInput]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (KernelInput -> Bool) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Bool
isDest (VName -> Bool) -> (KernelInput -> VName) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelInput -> VName
kernelInputArray) [KernelInput]
inputs

  Certificates -> BinderT (Rep m) m () -> BinderT (Rep m) m ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (BinderT (Rep m) m () -> BinderT (Rep m) m ())
-> BinderT (Rep m) m () -> BinderT (Rep m) m ()
forall a b. (a -> b) -> a -> b
$
    Seq (Stm (Rep m)) -> BinderT (Rep m) m ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms (Seq (Stm (Rep m)) -> BinderT (Rep m) m ())
-> BinderT (Rep m) m (Seq (Stm (Rep m))) -> BinderT (Rep m) m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Stm (Rep m) -> BinderT (Rep m) m (Stm (Rep m)))
-> Seq (Stm (Rep m)) -> BinderT (Rep m) m (Seq (Stm (Rep m)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Stm (Rep m) -> BinderT (Rep m) m (Stm (Rep m))
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Stm rep -> m (Stm rep)
renameStm
      (Seq (Stm (Rep m)) -> BinderT (Rep m) m (Seq (Stm (Rep m))))
-> BinderT (Rep m) m (Seq (Stm (Rep m)))
-> BinderT (Rep m) m (Seq (Stm (Rep m)))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SegOpLevel (Rep m)
-> Pattern (Rep m)
-> SubExp
-> [(VName, SubExp)]
-> [KernelInput]
-> [HistOp (Rep m)]
-> Lambda (Rep m)
-> [VName]
-> BinderT (Rep m) m (Seq (Stm (Rep m)))
forall rep (m :: * -> *).
(DistRep rep, MonadFreshNames m, HasScope rep m) =>
SegOpLevel rep
-> Pattern rep
-> SubExp
-> [(VName, SubExp)]
-> [KernelInput]
-> [HistOp rep]
-> Lambda rep
-> [VName]
-> m (Stms rep)
segHist SegOpLevel (Rep m)
lvl PatternT Type
Pattern (Rep m)
orig_pat SubExp
hist_w [(VName, SubExp)]
ispace [KernelInput]
inputs' [HistOp (Rep m)]
ops' Lambda (Rep m)
lam [VName]
arrs

determineReduceOp ::
  MonadBinder m =>
  Lambda SOACS ->
  [SubExp] ->
  m (Lambda SOACS, [SubExp], Shape)
determineReduceOp :: forall (m :: * -> *).
MonadBinder m =>
Lambda -> [SubExp] -> m (Lambda, [SubExp], ShapeBase SubExp)
determineReduceOp Lambda
lam [SubExp]
nes =
  -- FIXME? We are assuming that the accumulator is a replicate, and
  -- we fish out its value in a gross way.
  case (SubExp -> Maybe VName) -> [SubExp] -> Maybe [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> Maybe VName
subExpVar [SubExp]
nes of
    Just [VName]
ne_vs' -> do
      let (ShapeBase SubExp
shape, Lambda
lam') = Lambda -> (ShapeBase SubExp, Lambda)
isVectorMap Lambda
lam
      [SubExp]
nes' <- [VName] -> (VName -> m SubExp) -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
ne_vs' ((VName -> m SubExp) -> m [SubExp])
-> (VName -> m SubExp) -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
ne_v -> do
        Type
ne_v_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ne_v
        [Char] -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"hist_ne" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
            VName -> Slice SubExp -> BasicOp
Index VName
ne_v (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
              Type -> Slice SubExp -> Slice SubExp
fullSlice Type
ne_v_t (Slice SubExp -> Slice SubExp) -> Slice SubExp -> Slice SubExp
forall a b. (a -> b) -> a -> b
$
                Int -> DimIndex SubExp -> Slice SubExp
forall a. Int -> a -> [a]
replicate (ShapeBase SubExp -> Int
forall a. ArrayShape a => a -> Int
shapeRank ShapeBase SubExp
shape) (DimIndex SubExp -> Slice SubExp)
-> DimIndex SubExp -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
      (Lambda, [SubExp], ShapeBase SubExp)
-> m (Lambda, [SubExp], ShapeBase SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Lambda
lam', [SubExp]
nes', ShapeBase SubExp
shape)
    Maybe [VName]
Nothing ->
      (Lambda, [SubExp], ShapeBase SubExp)
-> m (Lambda, [SubExp], ShapeBase SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Lambda
lam, [SubExp]
nes, ShapeBase SubExp
forall a. Monoid a => a
mempty)

isVectorMap :: Lambda SOACS -> (Shape, Lambda SOACS)
isVectorMap :: Lambda -> (ShapeBase SubExp, Lambda)
isVectorMap Lambda
lam
  | [Let (Pattern [] [PatElemT (LetDec SOACS)]
pes) StmAux (ExpDec SOACS)
_ (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))] <-
      Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms SOACS -> [Stm SOACS]) -> Stms SOACS -> [Stm SOACS]
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda -> Body SOACS
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda
lam,
    Body SOACS -> [SubExp]
forall rep. BodyT rep -> [SubExp]
bodyResult (Lambda -> Body SOACS
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda
lam) [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (PatElemT Type -> SubExp) -> [PatElemT Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> (PatElemT Type -> VName) -> PatElemT Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT Type -> VName
forall dec. PatElemT dec -> VName
patElemName) [PatElemT Type]
[PatElemT (LetDec SOACS)]
pes,
    Just Lambda
map_lam <- ScremaForm SOACS -> Maybe Lambda
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form,
    [VName]
arrs [VName] -> [VName] -> Bool
forall a. Eq a => a -> a -> Bool
== (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName (Lambda -> [LParam SOACS]
forall rep. LambdaT rep -> [LParam rep]
lambdaParams Lambda
lam) =
    let (ShapeBase SubExp
shape, Lambda
lam') = Lambda -> (ShapeBase SubExp, Lambda)
isVectorMap Lambda
map_lam
     in ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
w] ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> ShapeBase SubExp
shape, Lambda
lam')
  | Bool
otherwise = (ShapeBase SubExp
forall a. Monoid a => a
mempty, Lambda
lam)

segmentedScanomapKernel ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  KernelNest ->
  [Int] ->
  SubExp ->
  Lambda rep ->
  Lambda rep ->
  [SubExp] ->
  [VName] ->
  DistNestT rep m (Maybe (Stms rep))
segmentedScanomapKernel :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> SubExp
-> Lambda rep
-> Lambda rep
-> [SubExp]
-> [VName]
-> DistNestT rep m (Maybe (Stms rep))
segmentedScanomapKernel KernelNest
nest [Int]
perm SubExp
segment_size Lambda rep
lam Lambda rep
map_lam [SubExp]
nes [VName]
arrs = do
  [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl <- (DistEnv rep m
 -> [SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep m (SegOpLevel rep))
-> DistNestT
     rep
     m
     ([SubExp]
      -> [Char]
      -> ThreadRecommendation
      -> BinderT rep m (SegOpLevel rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m
-> [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep m (SegOpLevel rep)
forall rep (m :: * -> *). DistEnv rep m -> MkSegLevel rep m
distSegLevel
  KernelNest
-> [Int]
-> Names
-> Names
-> [SubExp]
-> [VName]
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Names
-> Names
-> [SubExp]
-> [VName]
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
isSegmentedOp KernelNest
nest [Int]
perm (Lambda rep -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda rep
lam) (Lambda rep -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda rep
map_lam) [SubExp]
nes [] ((PatternT Type
  -> [(VName, SubExp)]
  -> [KernelInput]
  -> [SubExp]
  -> [VName]
  -> BinderT rep m ())
 -> DistNestT rep m (Maybe (Stms rep)))
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
forall a b. (a -> b) -> a -> b
$
    \PatternT Type
pat [(VName, SubExp)]
ispace [KernelInput]
inps [SubExp]
nes' [VName]
_ -> do
      let scan_op :: SegBinOp rep
scan_op = Commutativity
-> Lambda rep -> [SubExp] -> ShapeBase SubExp -> SegBinOp rep
forall rep.
Commutativity
-> Lambda rep -> [SubExp] -> ShapeBase SubExp -> SegBinOp rep
SegBinOp Commutativity
Noncommutative Lambda rep
lam [SubExp]
nes' ShapeBase SubExp
forall a. Monoid a => a
mempty
      SegOpLevel rep
lvl <- [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl (SubExp
segment_size SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
ispace) [Char]
"segscan" (ThreadRecommendation -> BinderT rep m (SegOpLevel rep))
-> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
forall a b. (a -> b) -> a -> b
$ SegVirt -> ThreadRecommendation
NoRecommendation SegVirt
SegNoVirt
      Stms rep -> BinderT rep m ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms (Stms rep -> BinderT rep m ())
-> BinderT rep m (Stms rep) -> BinderT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Stm rep -> BinderT rep m (Stm rep))
-> Stms rep -> BinderT rep m (Stms rep)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Stm rep -> BinderT rep m (Stm rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Stm rep -> m (Stm rep)
renameStm
        (Stms rep -> BinderT rep m (Stms rep))
-> BinderT rep m (Stms rep) -> BinderT rep m (Stms rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SegOpLevel rep
-> Pattern rep
-> SubExp
-> [SegBinOp rep]
-> Lambda rep
-> [VName]
-> [(VName, SubExp)]
-> [KernelInput]
-> BinderT rep m (Stms rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep, HasScope rep m) =>
SegOpLevel rep
-> Pattern rep
-> SubExp
-> [SegBinOp rep]
-> Lambda rep
-> [VName]
-> [(VName, SubExp)]
-> [KernelInput]
-> m (Stms rep)
segScan SegOpLevel rep
lvl PatternT Type
Pattern rep
pat SubExp
segment_size [SegBinOp rep
scan_op] Lambda rep
map_lam [VName]
arrs [(VName, SubExp)]
ispace [KernelInput]
inps

regularSegmentedRedomapKernel ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  KernelNest ->
  [Int] ->
  SubExp ->
  Commutativity ->
  Lambda rep ->
  Lambda rep ->
  [SubExp] ->
  [VName] ->
  DistNestT rep m (Maybe (Stms rep))
regularSegmentedRedomapKernel :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> SubExp
-> Commutativity
-> Lambda rep
-> Lambda rep
-> [SubExp]
-> [VName]
-> DistNestT rep m (Maybe (Stms rep))
regularSegmentedRedomapKernel KernelNest
nest [Int]
perm SubExp
segment_size Commutativity
comm Lambda rep
lam Lambda rep
map_lam [SubExp]
nes [VName]
arrs = do
  [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl <- (DistEnv rep m
 -> [SubExp]
 -> [Char]
 -> ThreadRecommendation
 -> BinderT rep m (SegOpLevel rep))
-> DistNestT
     rep
     m
     ([SubExp]
      -> [Char]
      -> ThreadRecommendation
      -> BinderT rep m (SegOpLevel rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks DistEnv rep m
-> [SubExp]
-> [Char]
-> ThreadRecommendation
-> BinderT rep m (SegOpLevel rep)
forall rep (m :: * -> *). DistEnv rep m -> MkSegLevel rep m
distSegLevel
  KernelNest
-> [Int]
-> Names
-> Names
-> [SubExp]
-> [VName]
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Names
-> Names
-> [SubExp]
-> [VName]
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
isSegmentedOp KernelNest
nest [Int]
perm (Lambda rep -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda rep
lam) (Lambda rep -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda rep
map_lam) [SubExp]
nes [] ((PatternT Type
  -> [(VName, SubExp)]
  -> [KernelInput]
  -> [SubExp]
  -> [VName]
  -> BinderT rep m ())
 -> DistNestT rep m (Maybe (Stms rep)))
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
forall a b. (a -> b) -> a -> b
$
    \PatternT Type
pat [(VName, SubExp)]
ispace [KernelInput]
inps [SubExp]
nes' [VName]
_ -> do
      let red_op :: SegBinOp rep
red_op = Commutativity
-> Lambda rep -> [SubExp] -> ShapeBase SubExp -> SegBinOp rep
forall rep.
Commutativity
-> Lambda rep -> [SubExp] -> ShapeBase SubExp -> SegBinOp rep
SegBinOp Commutativity
comm Lambda rep
lam [SubExp]
nes' ShapeBase SubExp
forall a. Monoid a => a
mempty
      SegOpLevel rep
lvl <- [SubExp]
-> [Char] -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
mk_lvl (SubExp
segment_size SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
ispace) [Char]
"segred" (ThreadRecommendation -> BinderT rep m (SegOpLevel rep))
-> ThreadRecommendation -> BinderT rep m (SegOpLevel rep)
forall a b. (a -> b) -> a -> b
$ SegVirt -> ThreadRecommendation
NoRecommendation SegVirt
SegNoVirt
      Stms rep -> BinderT rep m ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms (Stms rep -> BinderT rep m ())
-> BinderT rep m (Stms rep) -> BinderT rep m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Stm rep -> BinderT rep m (Stm rep))
-> Stms rep -> BinderT rep m (Stms rep)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Stm rep -> BinderT rep m (Stm rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Stm rep -> m (Stm rep)
renameStm
        (Stms rep -> BinderT rep m (Stms rep))
-> BinderT rep m (Stms rep) -> BinderT rep m (Stms rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SegOpLevel rep
-> Pattern rep
-> SubExp
-> [SegBinOp rep]
-> Lambda rep
-> [VName]
-> [(VName, SubExp)]
-> [KernelInput]
-> BinderT rep m (Stms rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep, HasScope rep m) =>
SegOpLevel rep
-> Pattern rep
-> SubExp
-> [SegBinOp rep]
-> Lambda rep
-> [VName]
-> [(VName, SubExp)]
-> [KernelInput]
-> m (Stms rep)
segRed SegOpLevel rep
lvl PatternT Type
Pattern rep
pat SubExp
segment_size [SegBinOp rep
red_op] Lambda rep
map_lam [VName]
arrs [(VName, SubExp)]
ispace [KernelInput]
inps

isSegmentedOp ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  KernelNest ->
  [Int] ->
  Names ->
  Names ->
  [SubExp] ->
  [VName] ->
  ( PatternT Type ->
    [(VName, SubExp)] ->
    [KernelInput] ->
    [SubExp] ->
    [VName] ->
    BinderT rep m ()
  ) ->
  DistNestT rep m (Maybe (Stms rep))
isSegmentedOp :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
KernelNest
-> [Int]
-> Names
-> Names
-> [SubExp]
-> [VName]
-> (PatternT Type
    -> [(VName, SubExp)]
    -> [KernelInput]
    -> [SubExp]
    -> [VName]
    -> BinderT rep m ())
-> DistNestT rep m (Maybe (Stms rep))
isSegmentedOp KernelNest
nest [Int]
perm Names
free_in_op Names
_free_in_fold_op [SubExp]
nes [VName]
arrs PatternT Type
-> [(VName, SubExp)]
-> [KernelInput]
-> [SubExp]
-> [VName]
-> BinderT rep m ()
m = MaybeT (DistNestT rep m) (Stms rep)
-> DistNestT rep m (Maybe (Stms rep))
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT (DistNestT rep m) (Stms rep)
 -> DistNestT rep m (Maybe (Stms rep)))
-> MaybeT (DistNestT rep m) (Stms rep)
-> DistNestT rep m (Maybe (Stms rep))
forall a b. (a -> b) -> a -> b
$ do
  -- We must verify that array inputs to the operation are inputs to
  -- the outermost loop nesting or free in the loop nest.  Nothing
  -- free in the op may be bound by the nest.  Furthermore, the
  -- neutral elements must be free in the loop nest.
  --
  -- We must summarise any names from free_in_op that are bound in the
  -- nest, and describe how to obtain them given segment indices.

  let bound_by_nest :: Names
bound_by_nest = KernelNest -> Names
boundInKernelNest KernelNest
nest

  ([(VName, SubExp)]
ispace, [KernelInput]
kernel_inps) <- KernelNest
-> MaybeT (DistNestT rep m) ([(VName, SubExp)], [KernelInput])
forall (m :: * -> *).
MonadFreshNames m =>
KernelNest -> m ([(VName, SubExp)], [KernelInput])
flatKernel KernelNest
nest

  Bool -> MaybeT (DistNestT rep m) () -> MaybeT (DistNestT rep m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Names
free_in_op Names -> Names -> Bool
`namesIntersect` Names
bound_by_nest) (MaybeT (DistNestT rep m) () -> MaybeT (DistNestT rep m) ())
-> MaybeT (DistNestT rep m) () -> MaybeT (DistNestT rep m) ()
forall a b. (a -> b) -> a -> b
$
    [Char] -> MaybeT (DistNestT rep m) ()
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"Non-fold lambda uses nest-bound parameters."

  let indices :: [VName]
indices = ((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst [(VName, SubExp)]
ispace

      prepareNe :: SubExp -> MaybeT (DistNestT rep m) SubExp
prepareNe (Var VName
v)
        | VName
v VName -> Names -> Bool
`nameIn` Names
bound_by_nest =
          [Char] -> MaybeT (DistNestT rep m) SubExp
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"Neutral element bound in nest"
      prepareNe SubExp
ne = SubExp -> MaybeT (DistNestT rep m) SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
ne

      prepareArr :: VName -> MaybeT (DistNestT rep m) (BinderT rep m VName)
prepareArr VName
arr =
        case (KernelInput -> Bool) -> [KernelInput] -> Maybe KernelInput
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
arr) (VName -> Bool) -> (KernelInput -> VName) -> KernelInput -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelInput -> VName
kernelInputName) [KernelInput]
kernel_inps of
          Just KernelInput
inp
            | KernelInput -> [SubExp]
kernelInputIndices KernelInput
inp [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
indices ->
              BinderT rep m VName
-> MaybeT (DistNestT rep m) (BinderT rep m VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (BinderT rep m VName
 -> MaybeT (DistNestT rep m) (BinderT rep m VName))
-> BinderT rep m VName
-> MaybeT (DistNestT rep m) (BinderT rep m VName)
forall a b. (a -> b) -> a -> b
$ VName -> BinderT rep m VName
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> BinderT rep m VName) -> VName -> BinderT rep m VName
forall a b. (a -> b) -> a -> b
$ KernelInput -> VName
kernelInputArray KernelInput
inp
          Maybe KernelInput
Nothing
            | Bool -> Bool
not (VName
arr VName -> Names -> Bool
`nameIn` Names
bound_by_nest) ->
              -- This input is something that is free inside
              -- the loop nesting. We will have to replicate
              -- it.
              BinderT rep m VName
-> MaybeT (DistNestT rep m) (BinderT rep m VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (BinderT rep m VName
 -> MaybeT (DistNestT rep m) (BinderT rep m VName))
-> BinderT rep m VName
-> MaybeT (DistNestT rep m) (BinderT rep m VName)
forall a b. (a -> b) -> a -> b
$
                [Char] -> Exp (Rep (BinderT rep m)) -> BinderT rep m VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Rep m) -> m VName
letExp
                  (VName -> [Char]
baseString VName
arr [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_repd")
                  (BasicOp -> ExpT rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT rep) -> BasicOp -> ExpT rep
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> ShapeBase SubExp) -> [SubExp] -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
ispace) (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
arr)
          Maybe KernelInput
_ ->
            [Char] -> MaybeT (DistNestT rep m) (BinderT rep m VName)
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"Input not free, perfectly mapped, or outermost."

  [SubExp]
nes' <- (SubExp -> MaybeT (DistNestT rep m) SubExp)
-> [SubExp] -> MaybeT (DistNestT rep m) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> MaybeT (DistNestT rep m) SubExp
prepareNe [SubExp]
nes

  [BinderT rep m VName]
mk_arrs <- (VName -> MaybeT (DistNestT rep m) (BinderT rep m VName))
-> [VName] -> MaybeT (DistNestT rep m) [BinderT rep m VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> MaybeT (DistNestT rep m) (BinderT rep m VName)
prepareArr [VName]
arrs

  DistNestT rep m (Stms rep) -> MaybeT (DistNestT rep m) (Stms rep)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (DistNestT rep m (Stms rep) -> MaybeT (DistNestT rep m) (Stms rep))
-> DistNestT rep m (Stms rep)
-> MaybeT (DistNestT rep m) (Stms rep)
forall a b. (a -> b) -> a -> b
$
    m (Stms rep) -> DistNestT rep m (Stms rep)
forall rep (m :: * -> *) a.
(LocalScope rep m, DistRep rep) =>
m a -> DistNestT rep m a
liftInner (m (Stms rep) -> DistNestT rep m (Stms rep))
-> m (Stms rep) -> DistNestT rep m (Stms rep)
forall a b. (a -> b) -> a -> b
$
      BinderT rep m () -> m (Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
BinderT rep m a -> m (Stms rep)
runBinderT'_ (BinderT rep m () -> m (Stms rep))
-> BinderT rep m () -> m (Stms rep)
forall a b. (a -> b) -> a -> b
$ do
        [VName]
nested_arrs <- [BinderT rep m VName] -> BinderT rep m [VName]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [BinderT rep m VName]
mk_arrs

        let pat :: PatternT Type
pat =
              [PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] ([PatElemT Type] -> PatternT Type)
-> [PatElemT Type] -> PatternT Type
forall a b. (a -> b) -> a -> b
$
                [Int] -> [PatElemT Type] -> [PatElemT Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([PatElemT Type] -> [PatElemT Type])
-> [PatElemT Type] -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$
                  PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements (PatternT Type -> [PatElemT Type])
-> PatternT Type -> [PatElemT Type]
forall a b. (a -> b) -> a -> b
$ LoopNesting -> PatternT Type
loopNestingPattern (LoopNesting -> PatternT Type) -> LoopNesting -> PatternT Type
forall a b. (a -> b) -> a -> b
$ KernelNest -> LoopNesting
forall a b. (a, b) -> a
fst KernelNest
nest

        PatternT Type
-> [(VName, SubExp)]
-> [KernelInput]
-> [SubExp]
-> [VName]
-> BinderT rep m ()
m PatternT Type
pat [(VName, SubExp)]
ispace [KernelInput]
kernel_inps [SubExp]
nes' [VName]
nested_arrs

permutationAndMissing :: PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing :: PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type])
permutationAndMissing PatternT Type
pat [SubExp]
res = do
  let pes :: [PatElemT Type]
pes = PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements PatternT Type
pat
      ([PatElemT Type]
_used, [PatElemT Type]
unused) =
        (PatElemT Type -> Bool)
-> [PatElemT Type] -> ([PatElemT Type], [PatElemT Type])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((VName -> Names -> Bool
`nameIn` [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
res) (VName -> Bool)
-> (PatElemT Type -> VName) -> PatElemT Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT Type -> VName
forall dec. PatElemT dec -> VName
patElemName) [PatElemT Type]
pes
      res_expanded :: [SubExp]
res_expanded = [SubExp]
res [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ (PatElemT Type -> SubExp) -> [PatElemT Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> (PatElemT Type -> VName) -> PatElemT Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT Type -> VName
forall dec. PatElemT dec -> VName
patElemName) [PatElemT Type]
unused
  [Int]
perm <- (PatElemT Type -> SubExp) -> [PatElemT Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> (PatElemT Type -> VName) -> PatElemT Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT Type -> VName
forall dec. PatElemT dec -> VName
patElemName) [PatElemT Type]
pes [SubExp] -> [SubExp] -> Maybe [Int]
forall a. Eq a => [a] -> [a] -> Maybe [Int]
`isPermutationOf` [SubExp]
res_expanded
  ([Int], [PatElemT Type]) -> Maybe ([Int], [PatElemT Type])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
perm, [PatElemT Type]
unused)

-- Add extra pattern elements to every kernel nesting level.
expandKernelNest ::
  MonadFreshNames m =>
  [PatElemT Type] ->
  KernelNest ->
  m KernelNest
expandKernelNest :: forall (m :: * -> *).
MonadFreshNames m =>
[PatElemT Type] -> KernelNest -> m KernelNest
expandKernelNest [PatElemT Type]
pes (LoopNesting
outer_nest, [LoopNesting]
inner_nests) = do
  let outer_size :: [SubExp]
outer_size =
        LoopNesting -> SubExp
loopNestingWidth LoopNesting
outer_nest SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:
        (LoopNesting -> SubExp) -> [LoopNesting] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map LoopNesting -> SubExp
loopNestingWidth [LoopNesting]
inner_nests
      inner_sizes :: [[SubExp]]
inner_sizes = [SubExp] -> [[SubExp]]
forall a. [a] -> [[a]]
tails ([SubExp] -> [[SubExp]]) -> [SubExp] -> [[SubExp]]
forall a b. (a -> b) -> a -> b
$ (LoopNesting -> SubExp) -> [LoopNesting] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map LoopNesting -> SubExp
loopNestingWidth [LoopNesting]
inner_nests
  LoopNesting
outer_nest' <- LoopNesting -> [SubExp] -> m LoopNesting
expandWith LoopNesting
outer_nest [SubExp]
outer_size
  [LoopNesting]
inner_nests' <- (LoopNesting -> [SubExp] -> m LoopNesting)
-> [LoopNesting] -> [[SubExp]] -> m [LoopNesting]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM LoopNesting -> [SubExp] -> m LoopNesting
expandWith [LoopNesting]
inner_nests [[SubExp]]
inner_sizes
  KernelNest -> m KernelNest
forall (m :: * -> *) a. Monad m => a -> m a
return (LoopNesting
outer_nest', [LoopNesting]
inner_nests')
  where
    expandWith :: LoopNesting -> [SubExp] -> m LoopNesting
expandWith LoopNesting
nest [SubExp]
dims = do
      [PatElemT Type]
pes' <- (PatElemT Type -> m (PatElemT Type))
-> [PatElemT Type] -> m [PatElemT Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([SubExp] -> PatElemT Type -> m (PatElemT Type)
forall {m :: * -> *} {dec}.
(MonadFreshNames m, Typed dec) =>
[SubExp] -> PatElemT dec -> m (PatElemT Type)
expandPatElemWith [SubExp]
dims) [PatElemT Type]
pes
      LoopNesting -> m LoopNesting
forall (m :: * -> *) a. Monad m => a -> m a
return
        LoopNesting
nest
          { loopNestingPattern :: PatternT Type
loopNestingPattern =
              [PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] ([PatElemT Type] -> PatternT Type)
-> [PatElemT Type] -> PatternT Type
forall a b. (a -> b) -> a -> b
$
                PatternT Type -> [PatElemT Type]
forall dec. PatternT dec -> [PatElemT dec]
patternElements (LoopNesting -> PatternT Type
loopNestingPattern LoopNesting
nest) [PatElemT Type] -> [PatElemT Type] -> [PatElemT Type]
forall a. Semigroup a => a -> a -> a
<> [PatElemT Type]
pes'
          }

    expandPatElemWith :: [SubExp] -> PatElemT dec -> m (PatElemT Type)
expandPatElemWith [SubExp]
dims PatElemT dec
pe = do
      VName
name <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> m VName) -> [Char] -> m VName
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT dec
pe
      PatElemT Type -> m (PatElemT Type)
forall (m :: * -> *) a. Monad m => a -> m a
return
        PatElemT dec
pe
          { patElemName :: VName
patElemName = VName
name,
            patElemDec :: Type
patElemDec = PatElemT dec -> Type
forall dec. Typed dec => PatElemT dec -> Type
patElemType PatElemT dec
pe Type -> ShapeBase SubExp -> Type
`arrayOfShape` [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp]
dims
          }

kernelOrNot ::
  (MonadFreshNames m, DistRep rep) =>
  Certificates ->
  Stm SOACS ->
  DistAcc rep ->
  PostStms rep ->
  DistAcc rep ->
  Maybe (Stms rep) ->
  DistNestT rep m (DistAcc rep)
kernelOrNot :: forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Certificates
-> Stm SOACS
-> DistAcc rep
-> PostStms rep
-> DistAcc rep
-> Maybe (Stms rep)
-> DistNestT rep m (DistAcc rep)
kernelOrNot Certificates
cs Stm SOACS
bnd DistAcc rep
acc PostStms rep
_ DistAcc rep
_ Maybe (Stms rep)
Nothing =
  Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
Stm SOACS -> DistAcc rep -> DistNestT rep m (DistAcc rep)
addStmToAcc (Certificates -> Stm SOACS -> Stm SOACS
forall rep. Certificates -> Stm rep -> Stm rep
certify Certificates
cs Stm SOACS
bnd) DistAcc rep
acc
kernelOrNot Certificates
cs Stm SOACS
_ DistAcc rep
_ PostStms rep
kernels DistAcc rep
acc' (Just Stms rep
bnds) = do
  PostStms rep -> DistNestT rep m ()
forall (m :: * -> *) rep.
Monad m =>
PostStms rep -> DistNestT rep m ()
addPostStms PostStms rep
kernels
  Stms rep -> DistNestT rep m ()
forall (m :: * -> *) rep. Monad m => Stms rep -> DistNestT rep m ()
postStm (Stms rep -> DistNestT rep m ()) -> Stms rep -> DistNestT rep m ()
forall a b. (a -> b) -> a -> b
$ (Stm rep -> Stm rep) -> Stms rep -> Stms rep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certificates -> Stm rep -> Stm rep
forall rep. Certificates -> Stm rep -> Stm rep
certify Certificates
cs) Stms rep
bnds
  DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a. Monad m => a -> m a
return DistAcc rep
acc'

distributeMap ::
  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
  MapLoop ->
  DistAcc rep ->
  DistNestT rep m (DistAcc rep)
distributeMap :: forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)
distributeMap (MapLoop Pattern
pat StmAux ()
aux SubExp
w Lambda
lam [VName]
arrs) DistAcc rep
acc =
  DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (DistAcc rep)
distribute
    (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PatternT Type
-> StmAux ()
-> SubExp
-> Lambda
-> [VName]
-> DistNestT rep m (DistAcc rep)
-> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, DistRep rep) =>
PatternT Type
-> StmAux ()
-> SubExp
-> Lambda
-> [VName]
-> DistNestT rep m (DistAcc rep)
-> DistNestT rep m (DistAcc rep)
mapNesting
      PatternT Type
Pattern
pat
      StmAux ()
aux
      SubExp
w
      Lambda
lam
      [VName]
arrs
      (DistAcc rep -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> DistNestT rep m (DistAcc rep)
distribute (DistAcc rep -> DistNestT rep m (DistAcc rep))
-> DistNestT rep m (DistAcc rep) -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, LocalScope rep m, DistRep rep) =>
DistAcc rep -> Stms SOACS -> DistNestT rep m (DistAcc rep)
distributeMapBodyStms DistAcc rep
acc' Stms SOACS
lam_bnds)
  where
    acc' :: DistAcc rep
acc' =
      DistAcc :: forall rep. Targets -> Stms rep -> DistAcc rep
DistAcc
        { distTargets :: Targets
distTargets =
            (PatternT Type, [SubExp]) -> Targets -> Targets
pushInnerTarget
              (PatternT Type
Pattern
pat, Body SOACS -> [SubExp]
forall rep. BodyT rep -> [SubExp]
bodyResult (Body SOACS -> [SubExp]) -> Body SOACS -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda -> Body SOACS
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda
lam)
              (Targets -> Targets) -> Targets -> Targets
forall a b. (a -> b) -> a -> b
$ DistAcc rep -> Targets
forall rep. DistAcc rep -> Targets
distTargets DistAcc rep
acc,
          distStms :: Stms rep
distStms = Stms rep
forall a. Monoid a => a
mempty
        }

    lam_bnds :: Stms SOACS
lam_bnds = Body SOACS -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda -> Body SOACS
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda
lam