{-# LANGUAGE TypeFamilies #-}

-- | This pass transforms parallelism expressed with arbitrarily nested SOACs to
-- instead be expressed with limited-nesting SegOps. This is the so-called
-- "flattening transformation" (sometimes called "vectorization", although we do
-- not use that term much in the Futhark compiler).
--
-- This is a sophisticated pass that does various clever things:
--
-- - Detects uniform nesting and flattens it more efficiently than the
--   nonuniform case.
--
-- - DPH-style vectorization avoidance.
--
-- - Incremental flattening ("Futhark.Pass.Flatten.Incremental").
--
-- - Intrablock flattening ("Futhark.Pass.Flatten.Intrablock").
--
-- The goal is that *any* Futhark program must be compilable parallel GPU code,
-- although in some cases the resulting code is not particularly efficient.
--
-- The idea is to perform distribution on one level at a time, and produce
-- "irregular maps" that can accept and produce irregular arrays. These
-- irregular maps will then be transformed into flat parallelism based on their
-- contents. If irregular maps contain only a single Stm, then it is fairly
-- straightforward, as we simply implement flattening rules for every single
-- kind of expression. Of course that is also somewhat inefficient, so we want
-- to support multiple Stms for things like scalar code.
--
-- Nomenclature:
--
-- A /map-nest/ is the collection of parallel operations enclosing some code. For
-- simplicity, we say "map-nest" even when the top level parallel operation is
-- actually a redomap or other screma.
--
-- An /irregular array/ is a multidimensional array like '[[1,2],[3]]', where rows
-- have different shapes. These are not directly supported in Futhark or in the
-- Futhark IR, but are encoded in various ways.
--
-- We say that an operation or type in a map-nest is /uniform/ when its size
-- (including internal sizes and sizes of inputs) and control flow is invariant
-- to the map-nest. Converse, it is /nonuniform/ when it is variant. When we
-- distribute a uniform statement, the intermediate results are regular, and
-- otherwise irregular. A statement that uses an irregular array is necessarily
-- nonuniform.
--
-- Take care not to confuse the terms "regular" and "uniform" - we say "regular"
-- only about arrays! "Uniform" is the general concept.
--
-- /Uniform nested parallelism/ is nested parallelism whose size is uniform to
-- the enclosing map nest, and which uses only variables whose types are
-- uniform, and which is enclosed in uniform control flow. /Nonuniform nested
-- parallelism/ is the converse. Many of the optimisations here are about
-- detecting the uniform case. We previously often used the terms "regular
-- nested parallelism" and "irregular nested parallelism", but this is now
-- discouraged, as explained above.
module Futhark.Pass.Flatten (flattenSOACs) where

import Control.Monad
import Data.Bifunctor (second)
import Data.Foldable
import Data.List qualified as L
import Data.Map qualified as M
import Data.Set qualified as S
import Futhark.Analysis.Alias (analyseBody)
import Futhark.IR.Aliases (Aliases, bodyAliases)
import Futhark.IR.GPU
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Pass
import Futhark.Pass.Flatten.BasicOp
import Futhark.Pass.Flatten.Builtins
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Pass.Flatten.Incremental
import Futhark.Pass.Flatten.Loop
import Futhark.Pass.Flatten.Match
import Futhark.Pass.Flatten.PreProcess
import Futhark.Pass.Flatten.SOAC
import Futhark.Pass.Flatten.WithAcc
import Futhark.Tools
import Futhark.Transform.FirstOrderTransform qualified as FOT
import Futhark.Transform.Rename
import Futhark.Transform.ToGPU (soacsLambdaToGPU, soacsStmToGPU)
import Prelude hiding (div, quot, rem)

type FunSizeParams = Name -> S.Set Int

-- | The irregularity handling mode requested by a statement, defaulting to the
-- mode already in effect. @#[flattening(sequentialise_nonuniform)]@ asks that
-- nonuniform nested parallelism be sequentialised rather than flattened; see
-- 'SequentialiseIrregularAll'.
irregularityFor :: DistIrregularity -> StmAux a -> DistIrregularity
irregularityFor :: forall a. DistIrregularity -> StmAux a -> DistIrregularity
irregularityFor DistIrregularity
irreg StmAux a
aux
  | Name -> [Attr] -> Attr
AttrComp Name
"flattening" [Attr
"sequentialise_nonuniform"] Attr -> Attrs -> Bool
`inAttrs` StmAux a -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux a
aux =
      DistIrregularity
SequentialiseIrregularAll
  | Bool
otherwise = DistIrregularity
irreg

flattenOpsFor ::
  Attrs ->
  FunHasParallelism ->
  FunSizeParams ->
  DistIrregularity ->
  SegLevel ->
  FlattenOps
flattenOpsFor :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> DistIrregularity
-> SegLevel
-> FlattenOps
flattenOpsFor Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams DistIrregularity
irreg SegLevel
lvl =
  FlattenOps
    { flattenSegLevel :: SegLevel
flattenSegLevel = SegLevel
lvl,
      flattenIrregularity :: DistIrregularity
flattenIrregularity = DistIrregularity
irreg,
      flattenFunHasParallelism :: FunHasParallelism
flattenFunHasParallelism = FunHasParallelism
funHasParallelism,
      flattenDistStmWith :: FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv
flattenDistStmWith = FunSizeParams
-> FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv
transformDistStm FunSizeParams
funSizeParams,
      flattenScalarStmAt :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
flattenScalarStmAt = SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
transformScalarStm,
      flattenTopLevelStm :: Stm SOACS -> FlattenM ()
flattenTopLevelStm = Attrs
-> FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()
transformTopLevelStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams
    }

transformScalarStms ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Stms SOACS ->
  FlattenM DistEnv
transformScalarStms :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stms SOACS
-> FlattenM DistEnv
transformScalarStms SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps [DistResult]
distres Stms SOACS
stms = do
  let bound_in_batch :: Names
bound_in_batch = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (Stm SOACS -> [VName]) -> [Stm SOACS] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat Type -> [VName])
-> (Stm SOACS -> Pat Type) -> Stm SOACS -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm SOACS -> Pat Type
Stm SOACS -> Pat (LetDec SOACS)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat) ([Stm SOACS] -> [VName]) -> [Stm SOACS] -> [VName]
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms SOACS
stms
      allCerts :: Certs
allCerts = (Stm SOACS -> Certs) -> [Stm SOACS] -> Certs
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Stm SOACS
stm -> DistInputs -> StmAux () -> DistEnv -> Certs
forall a. DistInputs -> StmAux a -> DistEnv -> Certs
distCerts DistInputs
inps (Stm SOACS -> StmAux (ExpDec SOACS)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm SOACS
stm) DistEnv
env) (Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms SOACS
stms)
      certs :: Certs
certs = [VName] -> Certs
Certs ([VName] -> Certs) -> [VName] -> Certs
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Names -> Bool
`notNameIn` Names
bound_in_batch) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Certs -> [VName]
unCerts Certs
allCerts
  vs <- Certs -> FlattenM [VName] -> FlattenM [VName]
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
certs (FlattenM [VName] -> FlattenM [VName])
-> FlattenM [VName] -> FlattenM [VName]
forall a b. (a -> b) -> a -> b
$ Name -> Exp (Rep FlattenM) -> FlattenM [VName]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp Name
"scalar_dist" (Exp GPU -> FlattenM [VName])
-> ((Segments -> FlattenM Result) -> FlattenM (Exp GPU))
-> (Segments -> FlattenM Result)
-> FlattenM [VName]
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Exp GPU -> FlattenM (Exp GPU)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Exp rep -> m (Exp rep)
renameExp (Exp GPU -> FlattenM (Exp GPU))
-> ((Segments -> FlattenM Result) -> FlattenM (Exp GPU))
-> (Segments -> FlattenM Result)
-> FlattenM (Exp GPU)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> Segments -> (Segments -> FlattenM Result) -> FlattenM (Exp GPU)
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel -> f SubExp -> (f SubExp -> m Result) -> m (Exp GPU)
segMap SegLevel
lvl Segments
segments ((Segments -> FlattenM Result) -> FlattenM [VName])
-> (Segments -> FlattenM Result) -> FlattenM [VName]
forall a b. (a -> b) -> a -> b
$ \Segments
is -> do
    Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs Segments
segments DistEnv
env (Segments -> Segments
forall a. [a] -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Segments
is) DistInputs
inps
    Stms (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms (Stms (Rep FlattenM) -> FlattenM ())
-> Stms (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ (Stm SOACS -> Stm (Rep FlattenM))
-> Stms SOACS -> Stms (Rep FlattenM)
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm SOACS -> Stm (Rep FlattenM)
Stm SOACS -> Stm GPU
soacsStmToGPU Stms SOACS
stms
    Result -> FlattenM Result
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result -> FlattenM Result) -> Result -> FlattenM Result
forall a b. (a -> b) -> a -> b
$ Segments -> Result
subExpsRes (Segments -> Result) -> Segments -> Result
forall a b. (a -> b) -> a -> b
$ (DistResult -> SubExp) -> [DistResult] -> Segments
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp) -> (DistResult -> VName) -> DistResult -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistResult -> VName
distResName) [DistResult]
distres
  insertRepsM (zip (map distResTag distres) $ map Regular vs) env

transformScalarStm ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Stm SOACS ->
  FlattenM DistEnv
transformScalarStm :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
transformScalarStm SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Stm SOACS
stm =
  SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stms SOACS
-> FlattenM DistEnv
transformScalarStms SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm)

-- | Transform a top-level 'Screma' by treating it as the empty-'Segments' case
-- of a nested one: it is not enclosed in any map-nest, so there are no
-- segments, the mapped arrays are plain regular top-level values
-- ('DistInputFree'), and the results are necessarily regular.
transformTopLevelScrema ::
  Attrs ->
  FunHasParallelism ->
  FunSizeParams ->
  Pat Type ->
  StmAux () ->
  SubExp ->
  [VName] ->
  ScremaForm SOACS ->
  FlattenM ()
transformTopLevelScrema :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> FlattenM ()
transformTopLevelScrema Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs ScremaForm SOACS
form = do
  let irreg :: DistIrregularity
irreg = DistIrregularity -> StmAux () -> DistIrregularity
forall a. DistIrregularity -> StmAux a -> DistIrregularity
irregularityFor DistIrregularity
DistributeIrregular StmAux ()
aux
      ops :: FlattenOps
ops = Attrs
-> FunHasParallelism
-> FunSizeParams
-> DistIrregularity
-> SegLevel
-> FlattenOps
flattenOpsFor Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams DistIrregularity
irreg SegLevel
defaultSegLevel
  arr_ts <- (VName -> FlattenM Type) -> [VName] -> FlattenM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
arrs
  -- 'flattenScrema' may bind the names of the pattern it is given (some paths
  -- bind them directly, others only insert reps), so we pass it a fresh pattern
  -- and bind the real pattern names ourselves from the result.
  nested_pat <- renamePat pat
  let inps = (VName -> Type -> (VName, DistInput))
-> [VName] -> [Type] -> DistInputs
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\VName
arr Type
t -> (VName
arr, VName -> Type -> DistInput
DistInputFree VName
arr Type
t)) [VName]
arrs [Type]
arr_ts
      res =
        (Int -> PatElem Type -> DistResult)
-> [Int] -> [PatElem Type] -> [DistResult]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          (\Int
i PatElem Type
pe -> ResTag -> DistType -> VName -> DistResult
DistResult (Int -> ResTag
ResTag Int
i) (Segments -> Rank -> Type -> DistType
DistType [] (Int -> Rank
Rank Int
0) (PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem Type
pe)) (PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe))
          [Int
0 ..]
          (Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
nested_pat)
  env <- flattenScrema ops [] (DistEnv mempty) inps res (nested_pat, aux) (w, arrs, form)
  forM_ (zip (patNames pat) res) $ \(VName
pat_v, DistResult
r) ->
    case ResTag -> DistEnv -> ResRep
resVar (DistResult -> ResTag
distResTag DistResult
r) DistEnv
env of
      Regular VName
v ->
        [VName] -> Exp (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
pat_v] (Exp (Rep FlattenM) -> FlattenM ())
-> Exp (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
      Irregular IrregularRep
_ ->
        [Char] -> FlattenM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"transformTopLevelScrema: top-level result cannot be irregular"

liftArg :: SegLevel -> Segments -> SubExp -> DistInputs -> DistEnv -> (SubExp, Diet) -> FlattenM [(SubExp, Diet)]
liftArg :: SegLevel
-> Segments
-> SubExp
-> DistInputs
-> DistEnv
-> (SubExp, Diet)
-> FlattenM [(SubExp, Diet)]
liftArg SegLevel
lvl Segments
segments SubExp
w DistInputs
inps DistEnv
env (SubExp
se, Diet
d) = do
  (_, rep) <- SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> SubExp
-> FlattenM (Type, ResRep)
liftSubExp SegLevel
lvl Segments
segments DistInputs
inps DistEnv
env SubExp
se
  case rep of
    Regular VName
v -> do
      v_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
      v' <-
        if arrayShape v_t == Shape [w]
          then pure v
          else
            letExp "lifted_arg_flat" . BasicOp $
              Reshape v $
                reshapeAll (arrayShape v_t) (Shape [w])
      pure [(Var v', d)]
    Irregular IrregularRep
irreg -> do
      vs <- SubExp -> IrregularRep -> FlattenM [VName]
irregularRepToFlatArrs SubExp
w IrregularRep
irreg
      -- Only apply the original diet to the 'elems' array.
      pure $ zip (map Var vs) $ replicate 4 Observe ++ [d]

liftRegArg :: SegLevel -> Segments -> SubExp -> DistInputs -> DistEnv -> (SubExp, Diet) -> FlattenM (SubExp, Diet)
liftRegArg :: SegLevel
-> Segments
-> SubExp
-> DistInputs
-> DistEnv
-> (SubExp, Diet)
-> FlattenM (SubExp, Diet)
liftRegArg SegLevel
lvl Segments
_segments SubExp
w DistInputs
inps DistEnv
env (SubExp
se, Diet
d) = do
  se_t <- DistInputs -> SubExp -> FlattenM Type
subExpInputType DistInputs
inps SubExp
se
  let se_shape = Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
se_t
      expected_shape = Segments -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
se_shape
  v <- liftSubExpRegular lvl [w] inps env expected_shape se
  pure (Var v, d)

-- Lifts a functions return type such that it matches the lifted functions
-- return type.
--
-- A lifted function corresponds to 'map f', which always produces fresh arrays.
-- No result therefore aliases anything - neither an argument nor another
-- result - which is recorded by giving each an empty 'RetAls'; see 'noRetAls'.
-- Maintaining this invariant may require inserting copies in the function body;
-- see 'freshenResult'.
liftRetType :: SubExp -> [RetType SOACS] -> [RetType GPU]
liftRetType :: SubExp -> [RetType SOACS] -> [RetType GPU]
liftRetType SubExp
w = [[ExtType]] -> [ExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[ExtType]] -> [ExtType])
-> ([ExtType] -> [[ExtType]]) -> [ExtType] -> [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, [[ExtType]]) -> [[ExtType]]
forall a b. (a, b) -> b
snd ((Int, [[ExtType]]) -> [[ExtType]])
-> ([ExtType] -> (Int, [[ExtType]])) -> [ExtType] -> [[ExtType]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> ExtType -> (Int, [ExtType]))
-> Int -> [ExtType] -> (Int, [[ExtType]])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL Int -> ExtType -> (Int, [ExtType])
liftType Int
0
  where
    liftType :: Int -> ExtType -> (Int, [ExtType])
liftType Int
i ExtType
rettype =
      let lifted :: [ExtType]
lifted = case ExtType
rettype of
            Prim PrimType
pt -> ExtType -> [ExtType]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExtType -> [ExtType]) -> ExtType -> [ExtType]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 6)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 6)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w]) NoMode
NoMode
            Array PrimType
pt ShapeBase (Ext SubExp)
_ NoMode
_ ->
              let num_data :: TypeBase shape o
num_data = PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
                  segs :: ExtType
segs = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 10)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 10)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w]) NoMode
NoMode
                  flags :: ExtType
flags = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 9)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 9)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
i]) NoMode
NoMode
                  offsets :: ExtType
offsets = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 8)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 8)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w]) NoMode
NoMode
                  elems :: ExtType
elems = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 7)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 7)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
i]) NoMode
NoMode
               in [ExtType
forall {shape} {o}. TypeBase shape o
num_data, ExtType
segs, ExtType
flags, ExtType
offsets, ExtType
elems]
            Acc {} -> [Char] -> [ExtType]
forall a. HasCallStack => [Char] -> a
error [Char]
"liftRetType: Acc"
            Mem {} -> [Char] -> [ExtType]
forall a. HasCallStack => [Char] -> a
error [Char]
"liftRetType: Mem"
       in (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [ExtType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtType]
lifted, [ExtType]
lifted)

liftRegularRetType :: DistInputs -> SubExp -> [RetType SOACS] -> [RetType GPU]
liftRegularRetType :: DistInputs -> SubExp -> [RetType SOACS] -> [RetType GPU]
liftRegularRetType DistInputs
inps SubExp
w = [[ExtType]] -> [ExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[ExtType]] -> [ExtType])
-> ([ExtType] -> [[ExtType]]) -> [ExtType] -> [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, [[ExtType]]) -> [[ExtType]]
forall a b. (a, b) -> b
snd ((Int, [[ExtType]]) -> [[ExtType]])
-> ([ExtType] -> (Int, [[ExtType]])) -> [ExtType] -> [[ExtType]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> ExtType -> (Int, [ExtType]))
-> Int -> [ExtType] -> (Int, [[ExtType]])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL Int -> ExtType -> (Int, [ExtType])
liftType Int
0
  where
    liftType :: Int -> ExtType -> (Int, [ExtType])
liftType Int
i ExtType
rettype =
      let lifted :: [ExtType]
lifted = case ExtType
rettype of
            Prim PrimType
pt -> ExtType -> [ExtType]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExtType -> [ExtType]) -> ExtType -> [ExtType]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 0)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 0)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w]) NoMode
NoMode
            Array PrimType
pt ShapeBase (Ext SubExp)
shape NoMode
_ ->
              if DistInputs -> RetType SOACS -> Bool
needsIrregularRetType DistInputs
inps ExtType
RetType SOACS
rettype
                then
                  let num_data :: TypeBase shape o
num_data = PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
                      segs :: ExtType
segs = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 4)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 4)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w]) NoMode
NoMode
                      flags :: ExtType
flags = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 3)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 3)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
i]) NoMode
NoMode
                      offsets :: ExtType
offsets = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 2)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 2)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w]) NoMode
NoMode
                      elems :: ExtType
elems = TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 1)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 1)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
i]) NoMode
NoMode
                   in [ExtType
forall {shape} {o}. TypeBase shape o
num_data, ExtType
segs, ExtType
flags, ExtType
offsets, ExtType
elems]
                else
                  ExtType -> [ExtType]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExtType -> [ExtType]) -> ExtType -> [ExtType]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 5)
-> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase (ShapeBase (Ext SubExp)) (ZonkAny 5)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape [SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w] ShapeBase (Ext SubExp)
-> ShapeBase (Ext SubExp) -> ShapeBase (Ext SubExp)
forall a. Semigroup a => a -> a -> a
<> ShapeBase (Ext SubExp)
shape) NoMode
NoMode
            Acc {} -> [Char] -> [ExtType]
forall a. HasCallStack => [Char] -> a
error [Char]
"liftRetType: Acc"
            Mem {} -> [Char] -> [ExtType]
forall a. HasCallStack => [Char] -> a
error [Char]
"liftRetType: Mem"
       in (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [ExtType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtType]
lifted, [ExtType]
lifted)

liftFunName :: Name -> Name
liftFunName :: Name -> Name
liftFunName Name
name = Name
name Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_lifted"

liftUniformFunName :: Name -> Name
liftUniformFunName :: Name -> Name
liftUniformFunName Name
name = Name
name Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_uniform_lifted"

flattenApply ::
  FunSizeParams ->
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  (Pat Type, StmAux ()) ->
  (Name, [(SubExp, Diet)], [(RetType SOACS, RetAls)], Safety) ->
  FlattenM DistEnv
flattenApply :: FunSizeParams
-> SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> (Pat Type, StmAux ())
-> (Name, [(SubExp, Diet)], [(RetType SOACS, RetAls)], Safety)
-> FlattenM DistEnv
flattenApply FunSizeParams
funSizeParams SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
pat, StmAux ()
aux) (Name
name, [(SubExp, Diet)]
args, [(RetType SOACS, RetAls)]
rettype, Safety
s) =
  case SegLevel
lvl of
    SegThread {} -> do
      let size_positions :: Set Int
size_positions = FunSizeParams
funSizeParams Name
name
          indexed_args :: [(Int, (SubExp, Diet))]
indexed_args = [Int] -> [(SubExp, Diet)] -> [(Int, (SubExp, Diet))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [(SubExp, Diet)]
args
          isSizeArg :: (Int, (SubExp, Diet)) -> Bool
isSizeArg = (Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Int
size_positions) (Int -> Bool)
-> ((Int, (SubExp, Diet)) -> Int) -> (Int, (SubExp, Diet)) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, (SubExp, Diet)) -> Int
forall a b. (a, b) -> a
fst
          ([(Int, (SubExp, Diet))]
size_args, [(Int, (SubExp, Diet))]
value_args) = ((Int, (SubExp, Diet)) -> Bool)
-> [(Int, (SubExp, Diet))]
-> ([(Int, (SubExp, Diet))], [(Int, (SubExp, Diet))])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (Int, (SubExp, Diet)) -> Bool
isSizeArg [(Int, (SubExp, Diet))]
indexed_args
      let nonuniform :: Bool
nonuniform = ((Int, (SubExp, Diet)) -> Bool) -> [(Int, (SubExp, Diet))] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (SubExp -> Bool)
-> ((Int, (SubExp, Diet)) -> SubExp)
-> (Int, (SubExp, Diet))
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst ((SubExp, Diet) -> SubExp)
-> ((Int, (SubExp, Diet)) -> (SubExp, Diet))
-> (Int, (SubExp, Diet))
-> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, (SubExp, Diet)) -> (SubExp, Diet)
forall a b. (a, b) -> b
snd) [(Int, (SubExp, Diet))]
size_args
          name' :: Name
name' = if Bool
nonuniform then Name -> Name
liftFunName Name
name else Name -> Name
liftUniformFunName Name
name
          mode :: LiftMode
mode = if Bool
nonuniform then LiftMode
NonUniformLift else LiftMode
UniformLift
      Name -> LiftMode -> FlattenM ()
demandLifted Name
name LiftMode
mode
      w <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"num_segments" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp (Segments -> TPrimExp Int64 VName
segmentCount Segments
segments)

      args' <-
        if nonuniform
          then
            ((w, Observe) :) . concat <$> mapM (liftArg lvl segments w inps env) args
          else do
            value_args' <- mapM (liftRegArg lvl segments w inps env . snd) value_args
            -- We do not lift 'size_args' because they correspond to size
            -- parameters, which are invariant in the uniform case.
            pure $ (w, Observe) : map snd size_args <> value_args'
      let rettype' =
            if Bool
nonuniform
              then [RetType GPU] -> [(RetType GPU, RetAls)]
noRetAls ([RetType GPU] -> [(RetType GPU, RetAls)])
-> [RetType GPU] -> [(RetType GPU, RetAls)]
forall a b. (a -> b) -> a -> b
$ SubExp -> [RetType SOACS] -> [RetType GPU]
liftRetType SubExp
w ([RetType SOACS] -> [RetType GPU])
-> [RetType SOACS] -> [RetType GPU]
forall a b. (a -> b) -> a -> b
$ ((ExtType, RetAls) -> RetType SOACS)
-> [(ExtType, RetAls)] -> [RetType SOACS]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
(ExtType, RetAls) -> RetType SOACS
forall a b. (a, b) -> a
fst [(ExtType, RetAls)]
[(RetType SOACS, RetAls)]
rettype
              else [RetType GPU] -> [(RetType GPU, RetAls)]
noRetAls ([RetType GPU] -> [(RetType GPU, RetAls)])
-> [RetType GPU] -> [(RetType GPU, RetAls)]
forall a b. (a -> b) -> a -> b
$ DistInputs -> SubExp -> [RetType SOACS] -> [RetType GPU]
liftRegularRetType DistInputs
inps SubExp
w ([RetType SOACS] -> [RetType GPU])
-> [RetType SOACS] -> [RetType GPU]
forall a b. (a -> b) -> a -> b
$ ((ExtType, RetAls) -> RetType SOACS)
-> [(ExtType, RetAls)] -> [RetType SOACS]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
(ExtType, RetAls) -> RetType SOACS
forall a b. (a, b) -> a
fst [(ExtType, RetAls)]
[(RetType SOACS, RetAls)]
rettype
      result <- letTupExp (name' <> "_res") $ Apply name' args' rettype' s
      let reps =
            if Bool
nonuniform
              then [ExtType] -> [VName] -> [ResRep]
forall s o. [TypeBase s o] -> [VName] -> [ResRep]
resultToResReps (((ExtType, RetAls) -> ExtType) -> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
forall a b. (a, b) -> a
fst [(ExtType, RetAls)]
[(RetType SOACS, RetAls)]
rettype) [VName]
result
              -- XXX: This could instead distinguish between regular and
              -- irregular results based on their return types.
              else [DistResult] -> [VName] -> [ResRep]
resultToResRepsByDistResult [DistResult]
res [VName]
result
      reps' <- zipWithM (reshapeLiftedApplyResult segments) (map fst rettype) reps
      insertRepsM (zip (map distResTag res) reps') env
    -- TODO: we currently do not handle intrablock function applications. It
    -- is possible we could do intrablock-level lifting of functions, but
    -- for now, we simply do not generate intrablock kernels if they would
    -- contain calls to parallel functions.
    SegLevel
_ ->
      if (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
res
        then SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
transformScalarStm SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Stm SOACS -> FlattenM DistEnv) -> Stm SOACS -> FlattenM DistEnv
forall a b. (a -> b) -> a -> b
$ Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux (Name
-> [(SubExp, Diet)]
-> [(RetType SOACS, RetAls)]
-> Safety
-> Exp SOACS
forall rep.
Name
-> [(SubExp, Diet)] -> [(RetType rep, RetAls)] -> Safety -> Exp rep
Apply Name
name [(SubExp, Diet)]
args [(RetType SOACS, RetAls)]
rettype Safety
s)
        else [Char] -> FlattenM DistEnv
forall a. HasCallStack => [Char] -> a
error [Char]
"Unhandled Apply in non SegThread Seglevel"

transformDistStm :: FunSizeParams -> FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv
transformDistStm :: FunSizeParams
-> FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv
transformDistStm FunSizeParams
_ FlattenOps
outer_ops Segments
segments DistEnv
env (DistStm DistInputs
inps [DistResult]
res (ScalarStm Stms SOACS
stms)) =
  SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stms SOACS
-> FlattenM DistEnv
transformScalarStms (FlattenOps -> SegLevel
flattenSegLevel FlattenOps
outer_ops) Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Stms SOACS
stms
transformDistStm FunSizeParams
funSizeParams FlattenOps
outer_ops Segments
segments DistEnv
env (DistStm DistInputs
inps [DistResult]
res (ParallelStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e))) = do
  case Exp SOACS
e of
    BasicOp BasicOp
op -> do
      let ~[DistResult
res'] = [DistResult]
res
          ~[PatElem Type
pe] = Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec SOACS)
pat
      FlattenOps
-> Segments
-> DistEnv
-> (DistInputs, DistResult, PatElem Type, StmAux (), BasicOp)
-> FlattenM DistEnv
flattenBasicOp FlattenOps
ops Segments
segments DistEnv
env (DistInputs
inps, DistResult
res', PatElem Type
pe, StmAux ()
StmAux (ExpDec SOACS)
aux, BasicOp
op)
    Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form) ->
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> (Pat Type, StmAux ())
-> (SubExp, [VName], ScremaForm SOACS)
-> FlattenM DistEnv
flattenScrema FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
Pat (LetDec SOACS)
pat, StmAux ()
StmAux (ExpDec SOACS)
aux) (SubExp
w, [VName]
arrs, ScremaForm SOACS
form)
    Match Segments
scrutinees [Case (Body SOACS)]
cases Body SOACS
defaultCase MatchDec (BranchType SOACS)
rt ->
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> Segments
-> [Case (Body SOACS)]
-> Body SOACS
-> MatchDec ExtType
-> FlattenM DistEnv
flattenMatch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
StmAux (ExpDec SOACS)
aux Segments
scrutinees [Case (Body SOACS)]
cases Body SOACS
defaultCase MatchDec ExtType
MatchDec (BranchType SOACS)
rt
    Apply Name
name [(SubExp, Diet)]
args [(RetType SOACS, RetAls)]
rettype Safety
s ->
      FunSizeParams
-> SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> (Pat Type, StmAux ())
-> (Name, [(SubExp, Diet)], [(RetType SOACS, RetAls)], Safety)
-> FlattenM DistEnv
flattenApply FunSizeParams
funSizeParams SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
Pat (LetDec SOACS)
pat, StmAux ()
StmAux (ExpDec SOACS)
aux) (Name
name, [(SubExp, Diet)]
args, [(RetType SOACS, RetAls)]
rettype, Safety
s)
    Loop [(FParam SOACS, SubExp)]
merge (ForLoop VName
i IntType
it SubExp
n) Body SOACS
body ->
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> (Pat Type, StmAux ())
-> ([(Param DeclType, SubExp)], LoopForm, Body SOACS)
-> FlattenM DistEnv
flattenLoop FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
Pat (LetDec SOACS)
pat, StmAux ()
StmAux (ExpDec SOACS)
aux) ([(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
merge, VName -> IntType -> SubExp -> LoopForm
ForLoop VName
i IntType
it SubExp
n, Body SOACS
body)
    Loop [(FParam SOACS, SubExp)]
merge (WhileLoop VName
cond) Body SOACS
body -> do
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> (Pat Type, StmAux ())
-> ([(Param DeclType, SubExp)], LoopForm, Body SOACS)
-> FlattenM DistEnv
flattenLoop FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
Pat (LetDec SOACS)
pat, StmAux ()
StmAux (ExpDec SOACS)
aux) ([(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
merge, VName -> LoopForm
WhileLoop VName
cond, Body SOACS
body)
    WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
lam ->
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> [WithAccInput SOACS]
-> Lambda SOACS
-> FlattenM DistEnv
flattenWithAcc FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux [WithAccInput SOACS]
inputs Lambda SOACS
lam
    Op (Hist SubExp
w [VName]
hist_inputs [HistOp SOACS]
hist_ops Lambda SOACS
bucket_fun) ->
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> (Pat Type, StmAux ())
-> (SubExp, [VName], [HistOp SOACS], Lambda SOACS)
-> FlattenM DistEnv
flattenHist FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
Pat (LetDec SOACS)
pat, StmAux ()
StmAux (ExpDec SOACS)
aux) (SubExp
w, [VName]
hist_inputs, [HistOp SOACS]
hist_ops, Lambda SOACS
bucket_fun)
    Op (FlatMap SubExp
w [VName]
arrs ExtLambda SOACS
lam) ->
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> SubExp
-> [VName]
-> ExtLambda SOACS
-> FlattenM DistEnv
flattenFlatMapNested FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
StmAux (ExpDec SOACS)
aux SubExp
w [VName]
arrs ExtLambda SOACS
lam
    Op (Stream {}) -> [Char] -> FlattenM DistEnv
forall a. HasCallStack => [Char] -> a
error [Char]
"transformDistStm: Stream should have been removed"
    Op (JVP {}) -> [Char] -> FlattenM DistEnv
forall a. HasCallStack => [Char] -> a
error [Char]
"Unhandled JVP"
    Op (VJP {}) -> [Char] -> FlattenM DistEnv
forall a. HasCallStack => [Char] -> a
error [Char]
"Unhandled VJP"
    Op (WithVJP {}) -> [Char] -> FlattenM DistEnv
forall a. HasCallStack => [Char] -> a
error [Char]
"Unhandled WithVJP"
  where
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
outer_ops
    ops :: FlattenOps
ops =
      FlattenOps
outer_ops
        { flattenIrregularity =
            irregularityFor (flattenIrregularity outer_ops) aux
        }

reshapeLiftedApplyResult :: Segments -> RetType SOACS -> ResRep -> FlattenM ResRep
reshapeLiftedApplyResult :: Segments -> RetType SOACS -> ResRep -> FlattenM ResRep
reshapeLiftedApplyResult Segments
segments Prim {} (Regular VName
v) = do
  v_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
  let expectedShape = Segments -> Shape
segmentsShape Segments
segments
  v' <-
    if arrayShape v_t == expectedShape
      then pure v
      else
        letExp "lifted_apply_res" . BasicOp $
          Reshape v $
            reshapeAll (arrayShape v_t) expectedShape
  pure $ Regular v'
reshapeLiftedApplyResult Segments
_ RetType SOACS
_ ResRep
rep =
  ResRep -> FlattenM ResRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ResRep
rep

liftBody :: Attrs -> FunHasParallelism -> FunSizeParams -> SegLevel -> SubExp -> DistInputs -> DistEnv -> DistStms -> Result -> FlattenM Result
liftBody :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> SegLevel
-> SubExp
-> DistInputs
-> DistEnv
-> DistStms
-> Result
-> FlattenM Result
liftBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams SegLevel
lvl SubExp
w DistInputs
inputs DistEnv
env DistStms
dstms Result
result = do
  let segments :: Segments
segments = [SubExp
w]
      ops :: FlattenOps
ops = Attrs
-> FunHasParallelism
-> FunSizeParams
-> DistIrregularity
-> SegLevel
-> FlattenOps
flattenOpsFor Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams DistIrregularity
DistributeIrregular SegLevel
lvl
  env' <- (DistEnv -> DistStm -> FlattenM DistEnv)
-> DistEnv -> DistStms -> FlattenM DistEnv
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv
flattenDistStm FlattenOps
ops Segments
segments) DistEnv
env DistStms
dstms
  result' <- mapM (liftResult lvl segments inputs env') result
  pure $ concat result'

liftUniformFunBody :: Attrs -> FunHasParallelism -> FunSizeParams -> SegLevel -> SubExp -> DistInputs -> DistEnv -> DistStms -> [RetType SOACS] -> Result -> FlattenM Result
liftUniformFunBody :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> SegLevel
-> SubExp
-> DistInputs
-> DistEnv
-> DistStms
-> [RetType SOACS]
-> Result
-> FlattenM Result
liftUniformFunBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams SegLevel
lvl SubExp
w DistInputs
inputs DistEnv
env DistStms
dstms [RetType SOACS]
rettype Result
result = do
  let segments :: Segments
segments = [SubExp
w]
      ops :: FlattenOps
ops = Attrs
-> FunHasParallelism
-> FunSizeParams
-> DistIrregularity
-> SegLevel
-> FlattenOps
flattenOpsFor Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams DistIrregularity
DistributeIrregular SegLevel
lvl
  env' <- (DistEnv -> DistStm -> FlattenM DistEnv)
-> DistEnv -> DistStms -> FlattenM DistEnv
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv
flattenDistStm FlattenOps
ops Segments
segments) DistEnv
env DistStms
dstms
  concat <$> zipWithM (liftRegResult lvl segments w inputs env') rettype result

-- | A lifted function must return fresh, non-aliasing arrays (as it
-- corresponds to 'map f'; see 'liftRetType').  This is not
-- automatically the case: a result may alias a parameter (when a value
-- is passed straight through), or the same array may be returned in
-- multiple result positions (which happens for functions that return
-- the same value more than once).  For every such result we insert a
-- copy to re-establish the invariant.  Results that are already fresh
-- are left untouched, so no superfluous copies are inserted.
freshenResult :: [FParam GPU] -> FlattenM Result -> FlattenM Result
freshenResult :: [FParam GPU] -> FlattenM Result -> FlattenM Result
freshenResult [FParam GPU]
params FlattenM Result
m = do
  (result, stms) <- FlattenM Result -> FlattenM (Result, Stms (Rep FlattenM))
forall a. FlattenM a -> FlattenM (a, Stms (Rep FlattenM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms FlattenM Result
m
  addStms stms
  let param_names = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
paramName [Param DeclType]
[FParam GPU]
params
      -- Transitive aliases of each result, including aliases with
      -- parameters and other results.
      als = Body (Aliases GPU) -> [Names]
forall rep. Aliased rep => Body rep -> [Names]
bodyAliases (AliasTable -> GBody GPU SubExpRes -> Body (Aliases GPU)
forall rep res.
(AliasableRep rep, IsResult res) =>
AliasTable -> GBody rep res -> GBody (Aliases rep) res
analyseBody AliasTable
forall a. Monoid a => a
mempty (BodyDec GPU -> Stms GPU -> Result -> GBody GPU SubExpRes
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body () Stms GPU
stms Result
result) :: Body (Aliases GPU))
  reverse . snd <$> foldM freshen (param_names, []) (zip result als)
  where
    freshen :: (Names, Result) -> (SubExpRes, Names) -> m (Names, Result)
freshen (Names
taken, Result
acc) (SubExpRes Certs
cs (Var VName
v), Names
v_als) = do
      v_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
      case v_t of
        Array {}
          | Names
taken Names -> Names -> Bool
`namesIntersect` Names
v_als -> do
              v' <- Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"fresh_result" (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
              pure (taken, SubExpRes cs (Var v') : acc)
        Type
_ ->
          (Names, Result) -> m (Names, Result)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Names
taken Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
v_als, Certs -> SubExp -> SubExpRes
SubExpRes Certs
cs (VName -> SubExp
Var VName
v) SubExpRes -> Result -> Result
forall a. a -> [a] -> [a]
: Result
acc)
    freshen (Names
taken, Result
acc) (SubExpRes
res', Names
_) =
      (Names, Result) -> m (Names, Result)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Names
taken, SubExpRes
res' SubExpRes -> Result -> Result
forall a. a -> [a] -> [a]
: Result
acc)

analyseFunParallelism :: [FunDef SOACS] -> M.Map Name Bool
analyseFunParallelism :: [FunDef SOACS] -> Map Name Bool
analyseFunParallelism [FunDef SOACS]
funs =
  [(Name, Bool)] -> Map Name Bool
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName FunDef SOACS
fun, Set Name -> FunHasParallelism
hasParallelFun Set Name
forall a. Monoid a => a
mempty (FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName FunDef SOACS
fun)) | FunDef SOACS
fun <- [FunDef SOACS]
funs]
  where
    funsByName :: Map Name (FunDef SOACS)
funsByName =
      [(Name, FunDef SOACS)] -> Map Name (FunDef SOACS)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName FunDef SOACS
fun, FunDef SOACS
fun) | FunDef SOACS
fun <- [FunDef SOACS]
funs]
    hasParallelFun :: Set Name -> FunHasParallelism
hasParallelFun Set Name
seen Name
fname
      | FunHasParallelism
isBuiltInFunction Name
fname =
          Bool
False
      -- avoid cycles even thought it is impossible now
      | Name
fname Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Name
seen =
          Bool
False
      | Just FunDef SOACS
fun <- Name -> Map Name (FunDef SOACS) -> Maybe (FunDef SOACS)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (FunDef SOACS)
funsByName =
          (Stm SOACS -> Bool) -> Stms SOACS -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (FunHasParallelism -> Stm SOACS -> Bool
isParallelStm (Set Name -> FunHasParallelism
hasParallelFun (Name -> Set Name -> Set Name
forall a. Ord a => a -> Set a -> Set a
S.insert Name
fname Set Name
seen))) (Stms SOACS -> Bool) -> Stms SOACS -> Bool
forall a b. (a -> b) -> a -> b
$
            Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$
              FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fun
      | Bool
otherwise =
          [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error ([Char] -> Bool) -> [Char] -> Bool
forall a b. (a -> b) -> a -> b
$ [Char]
"analyseFunParallelism: unknown function " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
fname

analyseFunSizeParams :: [FunDef SOACS] -> M.Map Name (S.Set Int)
analyseFunSizeParams :: [FunDef SOACS] -> Map Name (Set Int)
analyseFunSizeParams = [(Name, Set Int)] -> Map Name (Set Int)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, Set Int)] -> Map Name (Set Int))
-> ([FunDef SOACS] -> [(Name, Set Int)])
-> [FunDef SOACS]
-> Map Name (Set Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FunDef SOACS -> (Name, Set Int))
-> [FunDef SOACS] -> [(Name, Set Int)]
forall a b. (a -> b) -> [a] -> [b]
map FunDef SOACS -> (Name, Set Int)
forall {rep} {a}.
(FreeIn (RetType rep), Typed (FParamInfo rep), Num a, Enum a,
 Ord a) =>
FunDef rep -> (Name, Set a)
analyse
  where
    analyse :: FunDef rep -> (Name, Set a)
analyse FunDef rep
fd =
      let fparams :: [FParam rep]
fparams = FunDef rep -> [FParam rep]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef rep
fd
          rettype :: [(RetType rep, RetAls)]
rettype = FunDef rep -> [(RetType rep, RetAls)]
forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType FunDef rep
fd
          size_names :: Names
size_names = ([Type], [RetType rep]) -> Names
forall a. FreeIn a => a -> Names
freeIn ((FParam rep -> Type) -> [FParam rep] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> Type
forall dec. Typed dec => Param dec -> Type
paramType [FParam rep]
fparams, ((RetType rep, RetAls) -> RetType rep)
-> [(RetType rep, RetAls)] -> [RetType rep]
forall a b. (a -> b) -> [a] -> [b]
map (RetType rep, RetAls) -> RetType rep
forall a b. (a, b) -> a
fst [(RetType rep, RetAls)]
rettype)
          isSizeParam :: FParam rep -> Bool
isSizeParam FParam rep
p = FParam rep -> VName
forall dec. Param dec -> VName
paramName FParam rep
p VName -> Names -> Bool
`nameIn` Names
size_names
          indexed_params :: [(a, FParam rep)]
indexed_params = [a] -> [FParam rep] -> [(a, FParam rep)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a
0 ..] [FParam rep]
fparams
          size_params :: [(a, FParam rep)]
size_params = ((a, FParam rep) -> Bool) -> [(a, FParam rep)] -> [(a, FParam rep)]
forall a. (a -> Bool) -> [a] -> [a]
filter (FParam rep -> Bool
isSizeParam (FParam rep -> Bool)
-> ((a, FParam rep) -> FParam rep) -> (a, FParam rep) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, FParam rep) -> FParam rep
forall a b. (a, b) -> b
snd) [(a, FParam rep)]
indexed_params
       in (FunDef rep -> Name
forall rep. FunDef rep -> Name
funDefName FunDef rep
fd, [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList ([a] -> Set a) -> [a] -> Set a
forall a b. (a -> b) -> a -> b
$ ((a, FParam rep) -> a) -> [(a, FParam rep)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, FParam rep) -> a
forall a b. (a, b) -> a
fst [(a, FParam rep)]
size_params)

-- | Every result of a lifted function is fresh (see 'liftRetType'), so
-- none of them alias a parameter or each other.
noRetAls :: [RetType GPU] -> [(RetType GPU, RetAls)]
noRetAls :: [RetType GPU] -> [(RetType GPU, RetAls)]
noRetAls = (ExtType -> (ExtType, RetAls)) -> [ExtType] -> [(ExtType, RetAls)]
forall a b. (a -> b) -> [a] -> [b]
map (,RetAls
forall a. Monoid a => a
mempty)

-- | Impose attributes on the statements of a function body. This is used to
-- impose attributes on top level statements in lifted functions.
imposeAttrsBody :: Attrs -> Body SOACS -> Body SOACS
imposeAttrsBody :: Attrs -> Body SOACS -> Body SOACS
imposeAttrsBody Attrs
attrs Body SOACS
body =
  Body SOACS
body {bodyStms = fmap (imposeAttrs attrs) (bodyStms body)}

liftFunDef ::
  Attrs ->
  FunHasParallelism ->
  FunSizeParams ->
  Scope SOACS ->
  FunDef SOACS ->
  PassM (FunDef GPU, S.Set DemandFn)
liftFunDef :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> Scope SOACS
-> FunDef SOACS
-> PassM (FunDef GPU, Set DemandFn)
liftFunDef Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Scope SOACS
const_scope FunDef SOACS
fd = do
  let FunDef
        { funDefBody :: forall rep. FunDef rep -> Body rep
funDefBody = Body SOACS
body,
          funDefParams :: forall rep. FunDef rep -> [FParam rep]
funDefParams = [FParam SOACS]
fparams,
          funDefRetType :: forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType = [(RetType SOACS, RetAls)]
rettype
        } = FunDef SOACS
fd
  wp <- Name -> DeclType -> PassM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"w" (DeclType -> PassM (Param DeclType))
-> DeclType -> PassM (Param DeclType)
forall a b. (a -> b) -> a -> b
$ PrimType -> DeclType
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  let w = VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
wp
  (fparams', reps) <- mapAndUnzipM (liftParam w) fparams
  let fparams'' = Param DeclType
wp Param DeclType -> [Param DeclType] -> [Param DeclType]
forall a. a -> [a] -> [a]
: [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
fparams'
  let inputs = do
        (p, i) <- [Param DeclType] -> [Int] -> [(Param DeclType, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam SOACS]
fparams [Int
0 ..]
        pure (paramName p, DistInput (ResTag i) (paramType p))
  let rettype' =
        [RetType GPU] -> [(RetType GPU, RetAls)]
noRetAls ([RetType GPU] -> [(RetType GPU, RetAls)])
-> [RetType GPU] -> [(RetType GPU, RetAls)]
forall a b. (a -> b) -> a -> b
$
          SubExp -> [RetType SOACS] -> [RetType GPU]
liftRetType SubExp
w (((ExtType, RetAls) -> ExtType) -> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
forall a b. (a, b) -> a
fst [(ExtType, RetAls)]
[(RetType SOACS, RetAls)]
rettype)
  let (inputs', dstms) =
        distributeBody DistributeIrregular funHasParallelism const_scope [Var (paramName wp)] inputs $
          imposeAttrsBody attrs body
      env = Map ResTag ResRep -> DistEnv
DistEnv (Map ResTag ResRep -> DistEnv) -> Map ResTag ResRep -> DistEnv
forall a b. (a -> b) -> a -> b
$ [(ResTag, ResRep)] -> Map ResTag ResRep
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(ResTag, ResRep)] -> Map ResTag ResRep)
-> [(ResTag, ResRep)] -> Map ResTag ResRep
forall a b. (a -> b) -> a -> b
$ [ResTag] -> [ResRep] -> [(ResTag, ResRep)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> ResTag) -> [Int] -> [ResTag]
forall a b. (a -> b) -> [a] -> [b]
map Int -> ResTag
ResTag [Int
0 ..]) [ResRep]
reps
  -- Lift the body of the function and get the results, inserting copies as
  -- necessary to ensure the results are fresh (see 'freshenResult').
  (body', needs) <-
    runFlattenM (castScope const_scope <> scopeOfFParams fparams'') $
      buildBody_ . freshenResult fparams'' $
        liftBody attrs funHasParallelism funSizeParams defaultSegLevel w inputs' env dstms $
          bodyResult body
  let name = Name -> Name
liftFunName (Name -> Name) -> Name -> Name
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName FunDef SOACS
fd
  pure
    ( fd
        { funDefName = name,
          funDefBody = body',
          funDefParams = fparams'',
          funDefRetType = rettype',
          -- We have to remove #[cpu_function] if present because this function
          -- really contains parallel code now.
          funDefAttrs = funDefAttrs fd `withoutAttrs` oneAttr "cpu_function"
        },
      needs
    )

-- Here we assume that every type size is invariant and therefore every input
-- array is regular. As a result, parameters that correspond to type sizes are
-- not lifted and are also not part of 'DistInput'.
-- A uniformly lifted function can still return irregular arrays. This happens
-- when it returns an array whose dimension size was created in the function
-- body. In other words, the array has an existential size.
liftUniformFunDef ::
  Attrs ->
  FunHasParallelism ->
  FunSizeParams ->
  Scope SOACS ->
  FunDef SOACS ->
  PassM (FunDef GPU, S.Set DemandFn)
liftUniformFunDef :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> Scope SOACS
-> FunDef SOACS
-> PassM (FunDef GPU, Set DemandFn)
liftUniformFunDef Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Scope SOACS
const_scope FunDef SOACS
fd = do
  let FunDef
        { funDefBody :: forall rep. FunDef rep -> Body rep
funDefBody = Body SOACS
body,
          funDefParams :: forall rep. FunDef rep -> [FParam rep]
funDefParams = [FParam SOACS]
fparams,
          funDefRetType :: forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType = [(RetType SOACS, RetAls)]
rettype
        } = FunDef SOACS
fd
  wp <- Name -> DeclType -> PassM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"w" (DeclType -> PassM (Param DeclType))
-> DeclType -> PassM (Param DeclType)
forall a b. (a -> b) -> a -> b
$ PrimType -> DeclType
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  let w = VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
wp
  let size_positions = FunSizeParams
funSizeParams FunSizeParams -> FunSizeParams
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName FunDef SOACS
fd
      isSizeParam = (Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Int
size_positions) (Int -> Bool)
-> ((Int, Param DeclType) -> Int) -> (Int, Param DeclType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Param DeclType) -> Int
forall a b. (a, b) -> a
fst
      (indexed_sizes, indexed_values) =
        L.partition isSizeParam $ zip [0 ..] fparams
      fparam_sizes = ((Int, Param DeclType) -> Param DeclType)
-> [(Int, Param DeclType)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Param DeclType) -> Param DeclType
forall a b. (a, b) -> b
snd [(Int, Param DeclType)]
indexed_sizes
      fparams_explicit = ((Int, Param DeclType) -> Param DeclType)
-> [(Int, Param DeclType)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Param DeclType) -> Param DeclType
forall a b. (a, b) -> b
snd [(Int, Param DeclType)]
indexed_values

  (fparams_explicit', value_reps) <- mapAndUnzipM (liftRegularParam w) fparams_explicit
  let fparams'' = Param DeclType
wp Param DeclType -> [Param DeclType] -> [Param DeclType]
forall a. a -> [a] -> [a]
: [Param DeclType]
fparam_sizes [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. Semigroup a => a -> a -> a
<> [Param DeclType]
fparams_explicit'
  let inputs = do
        (p, i) <- [Param DeclType] -> [Int] -> [(Param DeclType, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
fparams_explicit [Int
0 ..]
        pure (paramName p, DistInput (ResTag i) (paramType p))
  let (inputs', dstms) =
        distributeBody DistributeIrregular funHasParallelism (const_scope <> scopeOfFParams fparam_sizes) [Var (paramName wp)] inputs $
          imposeAttrsBody attrs body
      env = Map ResTag ResRep -> DistEnv
DistEnv (Map ResTag ResRep -> DistEnv) -> Map ResTag ResRep -> DistEnv
forall a b. (a -> b) -> a -> b
$ [(ResTag, ResRep)] -> Map ResTag ResRep
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(ResTag, ResRep)] -> Map ResTag ResRep)
-> [(ResTag, ResRep)] -> Map ResTag ResRep
forall a b. (a -> b) -> a -> b
$ [ResTag] -> [ResRep] -> [(ResTag, ResRep)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> ResTag) -> [Int] -> [ResTag]
forall a b. (a -> b) -> [a] -> [b]
map Int -> ResTag
ResTag [Int
0 ..]) [ResRep]
value_reps
      rettype' =
        [RetType GPU] -> [(RetType GPU, RetAls)]
noRetAls ([RetType GPU] -> [(RetType GPU, RetAls)])
-> [RetType GPU] -> [(RetType GPU, RetAls)]
forall a b. (a -> b) -> a -> b
$
          DistInputs -> SubExp -> [RetType SOACS] -> [RetType GPU]
liftRegularRetType DistInputs
inputs' SubExp
w (((ExtType, RetAls) -> ExtType) -> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
forall a b. (a, b) -> a
fst [(ExtType, RetAls)]
[(RetType SOACS, RetAls)]
rettype)
  -- Lift the body of the function and get the results, inserting copies as
  -- necessary to ensure the results are fresh (see 'freshenResult').
  (body', needs) <-
    runFlattenM (castScope const_scope <> scopeOfFParams fparams'') $
      buildBody_ . freshenResult fparams'' $
        -- XXX: I think function lifting makes it more important to classify invariant
        -- results in bodies. Function bodies can produce values that are
        -- invariant to the map-nest, but at this point there is no opportunity to
        -- hoist them out of the nest.

        liftUniformFunBody attrs funHasParallelism funSizeParams defaultSegLevel w inputs' env dstms (map fst rettype) $
          bodyResult body
  let name = Name -> Name
liftUniformFunName (Name -> Name) -> Name -> Name
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName FunDef SOACS
fd
  pure
    ( fd
        { funDefName = name,
          funDefBody = body',
          funDefParams = fparams'',
          funDefRetType = rettype',
          -- We have to remove #[cpu_function] if present because this function
          -- really contains parallel code now.
          funDefAttrs = funDefAttrs fd `withoutAttrs` oneAttr "cpu_function"
        },
      needs
    )

transformLambda :: Attrs -> FunHasParallelism -> FunSizeParams -> Lambda SOACS -> FlattenM (Lambda GPU)
transformLambda :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> Lambda SOACS
-> FlattenM (Lambda GPU)
transformLambda Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Lambda [LParam SOACS]
params [Type]
ret Body SOACS
body) = do
  body' <- Scope GPU
-> FlattenM (GBody GPU SubExpRes) -> FlattenM (GBody GPU SubExpRes)
forall a. Scope GPU -> FlattenM a -> FlattenM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (LParamInfo GPU)] -> Scope GPU
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [LParam SOACS]
[Param (LParamInfo GPU)]
params) (FlattenM (GBody GPU SubExpRes) -> FlattenM (GBody GPU SubExpRes))
-> FlattenM (GBody GPU SubExpRes) -> FlattenM (GBody GPU SubExpRes)
forall a b. (a -> b) -> a -> b
$ Attrs
-> FunHasParallelism
-> FunSizeParams
-> Body SOACS
-> FlattenM (GBody GPU SubExpRes)
transformBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Body SOACS
body
  pure $ Lambda params ret body'

transformStm :: Attrs -> FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()
transformStm :: Attrs
-> FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()
transformStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac))
  | Attr
"sequential_outer" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux = do
      scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
      stms <- runBuilderT_ (FOT.transformSOAC pat soac) (castScope scope)
      transformStms attrs funHasParallelism funSizeParams $ fmap (certify (stmAuxCerts aux)) stms
transformStm Attrs
_ FunHasParallelism
_ FunSizeParams
_ Stm SOACS
stm
  | 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 (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm (Rep FlattenM) -> FlattenM ())
-> Stm (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Stm GPU
soacsStmToGPU Stm SOACS
stm
transformStm Attrs
_ FunHasParallelism
_ FunSizeParams
_ (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op (Hist SubExp
w [VName]
arrs [HistOp SOACS]
ops Lambda SOACS
bucket_fun))) =
  Certs -> FlattenM () -> FlattenM ()
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux) (FlattenM () -> FlattenM ()) -> FlattenM () -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ do
    res <-
      SegLevel
-> Name
-> Segments
-> [HistOp SOACS]
-> Lambda GPU
-> [VName]
-> (Segments -> FlattenM ())
-> FlattenM [VName]
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> Segments
-> [HistOp SOACS]
-> Lambda GPU
-> [VName]
-> (Segments -> m ())
-> m [VName]
genUniformSegHist
        SegLevel
defaultSegLevel
        Name
"topLevelSegHist"
        [SubExp
w]
        [HistOp SOACS]
ops
        (Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
bucket_fun)
        [VName]
arrs
        (FlattenM () -> Segments -> FlattenM ()
forall a b. a -> b -> a
const (FlattenM () -> Segments -> FlattenM ())
-> FlattenM () -> Segments -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ () -> FlattenM ()
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
    forM_ (zip (patNames pat) res) $ \(VName
v, VName
v') ->
      [VName] -> Exp (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep FlattenM) -> FlattenM ())
-> Exp (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v'
transformStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form)))
  | ScremaForm SOACS -> Bool
shouldDissectForm ScremaForm SOACS
form =
      [Char] -> FlattenM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"transformStm: complex Screma survived preprocessing"
  | Bool
otherwise =
      Attrs
-> FunHasParallelism
-> FunSizeParams
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> FlattenM ()
transformTopLevelScrema Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux SubExp
w [VName]
arrs ScremaForm SOACS
form
transformStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op (FlatMap SubExp
w [VName]
arrs ExtLambda SOACS
lam))) =
  Certs -> FlattenM () -> FlattenM ()
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux) (FlattenM () -> FlattenM ()) -> FlattenM () -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ FlattenOps
-> Pat Type -> SubExp -> [VName] -> ExtLambda SOACS -> FlattenM ()
flattenFlatMap FlattenOps
ops Pat Type
Pat (LetDec SOACS)
pat SubExp
w [VName]
arrs ExtLambda SOACS
lam
  where
    irreg :: DistIrregularity
irreg = DistIrregularity -> StmAux () -> DistIrregularity
forall a. DistIrregularity -> StmAux a -> DistIrregularity
irregularityFor DistIrregularity
DistributeIrregular StmAux ()
StmAux (ExpDec SOACS)
aux
    ops :: FlattenOps
ops = Attrs
-> FunHasParallelism
-> FunSizeParams
-> DistIrregularity
-> SegLevel
-> FlattenOps
flattenOpsFor Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams DistIrregularity
irreg SegLevel
defaultSegLevel
transformStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Loop [(FParam SOACS, SubExp)]
params LoopForm
form Body SOACS
body)) =
  Scope GPU -> FlattenM () -> FlattenM ()
forall a. Scope GPU -> FlattenM a -> FlattenM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (LoopForm -> Scope GPU
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> [FParam GPU] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (((FParam GPU, SubExp) -> FParam GPU)
-> [(FParam GPU, SubExp)] -> [FParam GPU]
forall a b. (a -> b) -> [a] -> [b]
map (FParam GPU, SubExp) -> FParam GPU
forall a b. (a, b) -> a
fst [(FParam SOACS, SubExp)]
[(FParam GPU, SubExp)]
params)) (FlattenM () -> FlattenM ()) -> FlattenM () -> FlattenM ()
forall a b. (a -> b) -> a -> b
$
    Stm (Rep FlattenM) -> FlattenM ()
Stm GPU -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm GPU -> FlattenM ())
-> (GBody GPU SubExpRes -> Stm GPU)
-> GBody GPU SubExpRes
-> FlattenM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
Pat (LetDec GPU)
pat StmAux (ExpDec SOACS)
StmAux (ExpDec GPU)
aux (Exp GPU -> Stm GPU)
-> (GBody GPU SubExpRes -> Exp GPU)
-> GBody GPU SubExpRes
-> Stm GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(FParam GPU, SubExp)]
-> LoopForm -> GBody GPU SubExpRes -> Exp GPU
forall rep.
[(FParam rep, SubExp)] -> LoopForm -> Body rep -> Exp rep
Loop [(FParam SOACS, SubExp)]
[(FParam GPU, SubExp)]
params LoopForm
form (GBody GPU SubExpRes -> FlattenM ())
-> FlattenM (GBody GPU SubExpRes) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Attrs
-> FunHasParallelism
-> FunSizeParams
-> Body SOACS
-> FlattenM (GBody GPU SubExpRes)
transformBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Body SOACS
body
transformStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Match Segments
ses [Case (Body SOACS)]
cases Body SOACS
def_body MatchDec (BranchType SOACS)
ret)) =
  Stm (Rep FlattenM) -> FlattenM ()
Stm GPU -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm GPU -> FlattenM ())
-> (Exp GPU -> Stm GPU) -> Exp GPU -> FlattenM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
Pat (LetDec GPU)
pat StmAux (ExpDec SOACS)
StmAux (ExpDec GPU)
aux
    (Exp GPU -> FlattenM ()) -> FlattenM (Exp GPU) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Segments
-> [Case (GBody GPU SubExpRes)]
-> GBody GPU SubExpRes
-> MatchDec (BranchType GPU)
-> Exp GPU
forall rep.
Segments
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match Segments
ses ([Case (GBody GPU SubExpRes)]
 -> GBody GPU SubExpRes -> MatchDec ExtType -> Exp GPU)
-> FlattenM [Case (GBody GPU SubExpRes)]
-> FlattenM (GBody GPU SubExpRes -> MatchDec ExtType -> Exp GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Case (Body SOACS) -> FlattenM (Case (GBody GPU SubExpRes)))
-> [Case (Body SOACS)] -> FlattenM [Case (GBody GPU SubExpRes)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Case (Body SOACS) -> FlattenM (Case (GBody GPU SubExpRes))
onCase [Case (Body SOACS)]
cases FlattenM (GBody GPU SubExpRes -> MatchDec ExtType -> Exp GPU)
-> FlattenM (GBody GPU SubExpRes)
-> FlattenM (MatchDec ExtType -> Exp GPU)
forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Attrs
-> FunHasParallelism
-> FunSizeParams
-> Body SOACS
-> FlattenM (GBody GPU SubExpRes)
transformBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Body SOACS
def_body FlattenM (MatchDec ExtType -> Exp GPU)
-> FlattenM (MatchDec ExtType) -> FlattenM (Exp GPU)
forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MatchDec ExtType -> FlattenM (MatchDec ExtType)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MatchDec ExtType
MatchDec (BranchType SOACS)
ret)
  where
    onCase :: Case (Body SOACS) -> FlattenM (Case (GBody GPU SubExpRes))
onCase = (Body SOACS -> FlattenM (GBody GPU SubExpRes))
-> Case (Body SOACS) -> FlattenM (Case (GBody GPU SubExpRes))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Case a -> f (Case b)
traverse (Attrs
-> FunHasParallelism
-> FunSizeParams
-> Body SOACS
-> FlattenM (GBody GPU SubExpRes)
transformBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams)
transformStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
withacc_lam)) = do
  Stm (Rep FlattenM) -> FlattenM ()
Stm GPU -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm GPU -> FlattenM ())
-> (Lambda GPU -> Stm GPU) -> Lambda GPU -> FlattenM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
Pat (LetDec GPU)
pat StmAux (ExpDec SOACS)
StmAux (ExpDec GPU)
aux (Exp GPU -> Stm GPU)
-> (Lambda GPU -> Exp GPU) -> Lambda GPU -> Stm GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [WithAccInput GPU] -> Lambda GPU -> Exp GPU
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc ((WithAccInput SOACS -> WithAccInput GPU)
-> [WithAccInput SOACS] -> [WithAccInput GPU]
forall a b. (a -> b) -> [a] -> [b]
map WithAccInput SOACS -> WithAccInput GPU
forall {a} {b} {b}.
(a, b, Maybe (Lambda SOACS, b)) -> (a, b, Maybe (Lambda GPU, b))
onInput [WithAccInput SOACS]
inputs)
    (Lambda GPU -> FlattenM ()) -> FlattenM (Lambda GPU) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Attrs
-> FunHasParallelism
-> FunSizeParams
-> Lambda SOACS
-> FlattenM (Lambda GPU)
transformLambda Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Lambda SOACS
withacc_lam
  where
    onInput :: (a, b, Maybe (Lambda SOACS, b)) -> (a, b, Maybe (Lambda GPU, b))
onInput (a
shape, b
arrs, Maybe (Lambda SOACS, b)
Nothing) =
      (a
shape, b
arrs, Maybe (Lambda GPU, b)
forall a. Maybe a
Nothing)
    onInput (a
shape, b
arrs, Just (Lambda SOACS
lam, b
nes)) =
      (a
shape, b
arrs, (Lambda GPU, b) -> Maybe (Lambda GPU, b)
forall a. a -> Maybe a
Just (Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
lam, b
nes))
transformStm Attrs
_ FunHasParallelism
_ FunSizeParams
_ Stm SOACS
stm = Stm (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm (Rep FlattenM) -> FlattenM ())
-> Stm (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Stm GPU
soacsStmToGPU Stm SOACS
stm

-- | Transform a statement that is not enclosed in any map-nest, whether it
-- occurs in a function body or was synthesised by flattening. This is where the
-- attributes imposed on the pass (see 'flattenSOACs') are put on the statement.
-- Nested statements receive them through the usual attribute propagation.
transformTopLevelStm :: Attrs -> FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()
transformTopLevelStm :: Attrs
-> FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()
transformTopLevelStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams =
  Attrs
-> FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()
transformStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Stm SOACS -> FlattenM ())
-> (Stm SOACS -> Stm SOACS) -> Stm SOACS -> FlattenM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Attrs -> Stm SOACS -> Stm SOACS
imposeAttrs Attrs
attrs

transformStms :: Attrs -> FunHasParallelism -> FunSizeParams -> Stms SOACS -> FlattenM ()
transformStms :: Attrs
-> FunHasParallelism -> FunSizeParams -> Stms SOACS -> FlattenM ()
transformStms Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Stms SOACS
stms =
  Scope GPU -> FlattenM () -> FlattenM ()
forall a. Scope GPU -> FlattenM a -> FlattenM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Scope SOACS -> Scope GPU
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope SOACS -> Scope GPU) -> Scope SOACS -> Scope GPU
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
stms) (FlattenM () -> FlattenM ()) -> FlattenM () -> FlattenM ()
forall a b. (a -> b) -> a -> b
$
    Seq () -> ()
forall m. Monoid m => Seq m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold (Seq () -> ()) -> FlattenM (Seq ()) -> FlattenM ()
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Stm SOACS -> FlattenM ()) -> Stms SOACS -> FlattenM (Seq ())
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Seq a -> f (Seq b)
traverse (Attrs
-> FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()
transformTopLevelStm Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams) Stms SOACS
stms

transformBody :: Attrs -> FunHasParallelism -> FunSizeParams -> Body SOACS -> FlattenM (Body GPU)
transformBody :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> Body SOACS
-> FlattenM (GBody GPU SubExpRes)
transformBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams (Body () Stms SOACS
stms Result
res) = FlattenM Result -> FlattenM (GBody (Rep FlattenM) SubExpRes)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (FlattenM Result -> FlattenM (GBody (Rep FlattenM) SubExpRes))
-> FlattenM Result -> FlattenM (GBody (Rep FlattenM) SubExpRes)
forall a b. (a -> b) -> a -> b
$ do
  Attrs
-> FunHasParallelism -> FunSizeParams -> Stms SOACS -> FlattenM ()
transformStms Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Stms SOACS
stms
  Result -> FlattenM Result
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res

transformFunDef ::
  Attrs ->
  FunHasParallelism ->
  FunSizeParams ->
  Scope SOACS ->
  FunDef SOACS ->
  PassM (FunDef GPU, S.Set DemandFn)
transformFunDef :: Attrs
-> FunHasParallelism
-> FunSizeParams
-> Scope SOACS
-> FunDef SOACS
-> PassM (FunDef GPU, Set DemandFn)
transformFunDef Attrs
_ FunHasParallelism
_ FunSizeParams
_ Scope SOACS
consts_scope FunDef SOACS
fd
  | Attr
"cpu_function" Attr -> Attrs -> Bool
`inAttrs` FunDef SOACS -> Attrs
forall rep. FunDef rep -> Attrs
funDefAttrs FunDef SOACS
fd =
      (,Set DemandFn
forall a. Monoid a => a
mempty) (FunDef GPU -> (FunDef GPU, Set DemandFn))
-> PassM (FunDef GPU) -> PassM (FunDef GPU, Set DemandFn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Scope GPU -> FunDef SOACS -> PassM (FunDef GPU)
forall (m :: * -> *) torep.
(MonadFreshNames m, FirstOrderRep torep) =>
Scope torep -> FunDef SOACS -> m (FunDef torep)
FOT.transformFunDef (Scope SOACS -> Scope GPU
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope Scope SOACS
consts_scope) FunDef SOACS
fd
transformFunDef Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Scope SOACS
consts_scope FunDef SOACS
fd = do
  let FunDef
        { funDefBody :: forall rep. FunDef rep -> Body rep
funDefBody = Body SOACS
body,
          funDefParams :: forall rep. FunDef rep -> [FParam rep]
funDefParams = [FParam SOACS]
fparams,
          funDefRetType :: forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType = [(RetType SOACS, RetAls)]
rettype
        } = FunDef SOACS
fd
  (body', needs) <-
    Scope GPU
-> FlattenM (GBody GPU SubExpRes)
-> PassM (GBody GPU SubExpRes, Set DemandFn)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope GPU -> FlattenM a -> m (a, Set DemandFn)
runFlattenM ([Param DeclType] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param DeclType]
[FParam SOACS]
fparams Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope SOACS -> Scope GPU
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope Scope SOACS
consts_scope) (FlattenM (GBody GPU SubExpRes)
 -> PassM (GBody GPU SubExpRes, Set DemandFn))
-> FlattenM (GBody GPU SubExpRes)
-> PassM (GBody GPU SubExpRes, Set DemandFn)
forall a b. (a -> b) -> a -> b
$
      Attrs
-> FunHasParallelism
-> FunSizeParams
-> Body SOACS
-> FlattenM (GBody GPU SubExpRes)
transformBody Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Body SOACS
body
  pure
    ( fd
        { funDefBody = body',
          funDefRetType = rettype,
          funDefParams = fparams
        },
      needs
    )

liftUntilFixedPoint ::
  Prog SOACS ->
  Attrs ->
  FunHasParallelism ->
  FunSizeParams ->
  Scope SOACS ->
  S.Set DemandFn ->
  S.Set DemandFn ->
  PassM [FunDef GPU]
liftUntilFixedPoint :: Prog SOACS
-> Attrs
-> FunHasParallelism
-> FunSizeParams
-> Scope SOACS
-> Set DemandFn
-> Set DemandFn
-> PassM [FunDef GPU]
liftUntilFixedPoint Prog SOACS
prog Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Scope SOACS
consts_scope Set DemandFn
made Set DemandFn
needed = do
  let made' :: Set DemandFn
made' = Set DemandFn
made Set DemandFn -> Set DemandFn -> Set DemandFn
forall a. Semigroup a => a -> a -> a
<> Set DemandFn
needed
  (lifted_funs, new_needed) <-
    (([FunDef GPU], [Set DemandFn]) -> ([FunDef GPU], Set DemandFn))
-> PassM ([FunDef GPU], [Set DemandFn])
-> PassM ([FunDef GPU], Set DemandFn)
forall a b. (a -> b) -> PassM a -> PassM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Set DemandFn] -> Set DemandFn)
-> ([FunDef GPU], [Set DemandFn]) -> ([FunDef GPU], Set DemandFn)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Set DemandFn -> Set DemandFn -> Set DemandFn
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Set DemandFn
made') (Set DemandFn -> Set DemandFn)
-> ([Set DemandFn] -> Set DemandFn)
-> [Set DemandFn]
-> Set DemandFn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Set DemandFn] -> Set DemandFn
forall a. Monoid a => [a] -> a
mconcat)) (PassM ([FunDef GPU], [Set DemandFn])
 -> PassM ([FunDef GPU], Set DemandFn))
-> PassM ([FunDef GPU], [Set DemandFn])
-> PassM ([FunDef GPU], Set DemandFn)
forall a b. (a -> b) -> a -> b
$
      (DemandFn -> PassM (FunDef GPU, Set DemandFn))
-> [DemandFn] -> PassM ([FunDef GPU], [Set DemandFn])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM DemandFn -> PassM (FunDef GPU, Set DemandFn)
mkDemanded ([DemandFn] -> PassM ([FunDef GPU], [Set DemandFn]))
-> [DemandFn] -> PassM ([FunDef GPU], [Set DemandFn])
forall a b. (a -> b) -> a -> b
$
        Set DemandFn -> [DemandFn]
forall a. Set a -> [a]
S.toList Set DemandFn
needed
  if new_needed == mempty
    then pure lifted_funs
    else
      (lifted_funs ++)
        <$> liftUntilFixedPoint prog attrs funHasParallelism funSizeParams consts_scope made' new_needed
  where
    mkDemanded :: DemandFn -> PassM (FunDef GPU, Set DemandFn)
mkDemanded (DemandLifted Name
fname LiftMode
mode) =
      case (FunDef SOACS -> Bool) -> [FunDef SOACS] -> Maybe (FunDef SOACS)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Name -> FunHasParallelism
forall a. Eq a => a -> a -> Bool
== Name
fname) FunHasParallelism -> (FunDef SOACS -> Name) -> FunDef SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName) ([FunDef SOACS] -> Maybe (FunDef SOACS))
-> [FunDef SOACS] -> Maybe (FunDef SOACS)
forall a b. (a -> b) -> a -> b
$ Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog of
        Just FunDef SOACS
fundef ->
          case LiftMode
mode of
            LiftMode
UniformLift -> Attrs
-> FunHasParallelism
-> FunSizeParams
-> Scope SOACS
-> FunDef SOACS
-> PassM (FunDef GPU, Set DemandFn)
liftUniformFunDef Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Scope SOACS
consts_scope FunDef SOACS
fundef
            LiftMode
NonUniformLift -> Attrs
-> FunHasParallelism
-> FunSizeParams
-> Scope SOACS
-> FunDef SOACS
-> PassM (FunDef GPU, Set DemandFn)
liftFunDef Attrs
attrs FunHasParallelism
funHasParallelism FunSizeParams
funSizeParams Scope SOACS
consts_scope FunDef SOACS
fundef
        Maybe (FunDef SOACS)
Nothing -> [Char] -> PassM (FunDef GPU, Set DemandFn)
forall a. HasCallStack => [Char] -> a
error ([Char] -> PassM (FunDef GPU, Set DemandFn))
-> [Char] -> PassM (FunDef GPU, Set DemandFn)
forall a b. (a -> b) -> a -> b
$ [Char]
"mkDemanded: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
forall a. Show a => a -> [Char]
show Name
fname
    mkDemanded (DemandBuiltin BuiltinFn
b) = (FunDef GPU, Set DemandFn) -> PassM (FunDef GPU, Set DemandFn)
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BuiltinFn -> FunDef GPU
builtinFunDef BuiltinFn
b, Set DemandFn
forall a. Monoid a => a
mempty)

transformProg :: Attrs -> Prog SOACS -> PassM (Prog GPU)
transformProg :: Attrs -> Prog SOACS -> PassM (Prog GPU)
transformProg Attrs
attrs Prog SOACS
prog = do
  progAfterPreProcessing <- Prog SOACS -> PassM (Prog SOACS)
preprocessProg Prog SOACS
prog
  let consts = Prog SOACS -> Stms SOACS
forall rep. Prog rep -> Stms rep
progConsts Prog SOACS
progAfterPreProcessing
      consts_scope = Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
consts
      funs = Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
progAfterPreProcessing
      funParallelism = [FunDef SOACS] -> Map Name Bool
analyseFunParallelism [FunDef SOACS]
funs
      size_param_map = [FunDef SOACS] -> Map Name (Set Int)
analyseFunSizeParams [FunDef SOACS]
funs
      funHasParallelism Name
fname =
        Bool -> Name -> Map Name Bool -> Bool
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ FunHasParallelism
isBuiltInFunction Name
fname) Name
fname Map Name Bool
funParallelism
      funSizeParams Name
fname =
        Set Int -> Name -> Map Name (Set Int) -> Set Int
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Set Int
forall a. Monoid a => a
mempty Name
fname Map Name (Set Int)
size_param_map
  (consts', consts_needs) <-
    runFlattenM mempty $ collectStms_ $ transformStms attrs funHasParallelism funSizeParams consts
  (funs', funs_needs) <-
    second mconcat
      <$> mapAndUnzipM (transformFunDef attrs funHasParallelism funSizeParams consts_scope) funs

  -- Now do fixpoint iteration until all needed functions have been provided.
  lifted_funs <-
    liftUntilFixedPoint
      prog
      attrs
      funHasParallelism
      funSizeParams
      consts_scope
      mempty
      (consts_needs <> funs_needs)

  pure $
    prog
      { progConsts = consts',
        progFuns = lifted_funs <> funs'
      }

-- | Transform a SOACS program to a GPU program, using flattening.
--
-- Parameterised by the names of flattening attributes to impose on top-level
-- SOACs that do not carry flattening attributes of their own.
flattenSOACs :: [Name] -> Pass SOACS GPU
flattenSOACs :: [Name] -> Pass SOACS GPU
flattenSOACs [Name]
attr_names =
  Pass
    { passName :: [Char]
passName = [Char]
"flatten",
      passDescription :: [Char]
passDescription = [Char]
"Perform full flattening",
      passFunction :: Prog SOACS -> PassM (Prog GPU)
passFunction = Attrs -> Prog SOACS -> PassM (Prog GPU)
transformProg Attrs
attrs
    }
  where
    attrs :: Attrs
attrs =
      (Name -> Attrs) -> [Name] -> Attrs
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Attr -> Attrs
oneAttr (Attr -> Attrs) -> (Name -> Attr) -> Name -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Attr] -> Attr
AttrComp Name
"flattening" ([Attr] -> Attr) -> (Name -> [Attr]) -> Name -> Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Attr -> [Attr]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attr -> [Attr]) -> (Name -> Attr) -> Name -> [Attr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Attr
AttrName) [Name]
attr_names
{-# NOINLINE flattenSOACs #-}