{-# LANGUAGE LambdaCase #-}

-- | Flattening rules for SOACs.
module Futhark.Pass.Flatten.SOAC
  ( flattenScrema,
    flattenHist,
    flattenFlatMap,
    flattenFlatMapNested,
  )
where

import Control.Monad
import Data.Containers.ListUtils (nubOrd)
import Data.Foldable
import Data.Functor.Identity (runIdentity)
import Data.Map qualified as M
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Tuple.Solo
import Futhark.IR.GPU
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Pass.Flatten.Incremental
import Futhark.Pass.Flatten.Intrablock qualified as Intrablock
import Futhark.Pass.Flatten.PreProcess
import Futhark.Tools
import Futhark.Transform.FirstOrderTransform qualified as FOT
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Transform.ToGPU (soacsLambdaToGPU)
import Futhark.Util (mapAccumLM)
import Futhark.Util.IntegralExp
import Prelude hiding (div, quot, rem)

-- | How the results of a nested map are represented, determined by whether the
-- map width is uniform (invariant to the enclosing nest) or nonuniform; see
-- 'transformInnerMap'.
data InnerMapMode
  = -- | Uniform width: the results are regular arrays that keep the full
    -- multi-dimensional shape (the enclosing segments followed by the map
    -- width).
    MultiDim
  | -- | Nonuniform width: the results are irregular, flattened into a single
    -- segment dimension.
    SingleDim

freeWithTypeDeps :: DistInputs -> Names -> FlattenM [VName]
freeWithTypeDeps :: DistInputs -> Names -> FlattenM [VName]
freeWithTypeDeps DistInputs
inps Names
free = do
  let free_names :: [VName]
free_names = Names -> [VName]
namesToList Names
free
  free_sizes <- (Type -> Names) -> [Type] -> Names
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Type -> Names
forall a. FreeIn a => a -> Names
freeIn ([Type] -> Names) -> FlattenM [Type] -> FlattenM Names
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (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 (DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps) [VName]
free_names
  pure $ nubOrd $ namesToList free_sizes <> free_names

-- Reduction or scan operators may not have any free variables that are variant
-- to the nest (that is, are inputs to the distributed operation), and must
-- operate on primitive types. This is because we would be unable to express
-- them as SegScan/SegReds. Fixing this would require modifications to the
-- SegOp representation, but it is likely not worth it, as such operators are
-- extremely rare - and we can just fall back on sequentialising the SOAC and
-- flattening the resulting loop.
suitableOperator :: DistEnv -> DistInputs -> Lambda SOACS -> [SubExp] -> Bool
suitableOperator :: DistEnv -> DistInputs -> Lambda SOACS -> Segments -> Bool
suitableOperator DistEnv
_env DistInputs
inps Lambda SOACS
lam Segments
_nes =
  (VName -> Bool) -> Names -> Bool
allNames VName -> Bool
notVariant (Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
lam)
    Bool -> Bool -> Bool
&& (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam)
  where
    notVariant :: VName -> Bool
notVariant = Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (SubExp -> Bool) -> (VName -> SubExp) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var

suitableUniformOperator :: DistInputs -> Lambda SOACS -> Bool
suitableUniformOperator :: DistInputs -> Lambda SOACS -> Bool
suitableUniformOperator DistInputs
inps Lambda SOACS
lam =
  (VName -> Bool) -> Names -> Bool
allNames (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (SubExp -> Bool) -> (VName -> SubExp) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) (Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
lam)

regularToReplicatedIrregularRep ::
  SegLevel ->
  Segments ->
  VName ->
  VName ->
  FlattenM IrregularRep
regularToReplicatedIrregularRep :: SegLevel -> Segments -> VName -> VName -> FlattenM IrregularRep
regularToReplicatedIrregularRep SegLevel
lvl Segments
segments VName
ws_data VName
v' = do
  ws_prod <- Int -> Type -> SubExp
forall o. Int -> TypeBase (ShapeBase SubExp) o -> SubExp
arraySize Int
0 (Type -> SubExp) -> FlattenM Type -> FlattenM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ws_data
  arr_t <- lookupType v'
  segment_size <-
    letSubExp "reg_seg_size" <=< toExp . product . map pe64 $
      drop (segmentsRank segments) (arrayDims arr_t)
  num_elems <-
    letSubExp "reg_num_elems" <=< toExp $ product $ map pe64 $ arrayDims arr_t
  arr_D <-
    letExp "reg_D" . BasicOp $
      Reshape v' (reshapeAll (arrayShape arr_t) (Shape [num_elems]))
  arr_F <- letExp "reg_F" <=< segMap lvl (MkSolo num_elems) $ \(MkSolo SubExp
i) -> do
    flag <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"flag" (Exp GPU -> FlattenM SubExp)
-> (TPrimExp Bool VName -> FlattenM (Exp GPU))
-> TPrimExp Bool VName
-> FlattenM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< TPrimExp Bool VName -> FlattenM (Exp (Rep FlattenM))
TPrimExp Bool VName -> FlattenM (Exp GPU)
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Bool VName -> m (Exp (Rep m))
toExp (TPrimExp Bool VName -> FlattenM SubExp)
-> TPrimExp Bool VName -> FlattenM SubExp
forall a b. (a -> b) -> a -> b
$ (SubExp -> TPrimExp Int64 VName
pe64 SubExp
i TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`rem` SubExp -> TPrimExp Int64 VName
pe64 SubExp
segment_size) TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int64 VName
0
    pure [subExpRes flag]

  arr_S <-
    letExp "reg_segments" . BasicOp $
      Replicate (Shape [ws_prod]) segment_size
  arr_O <- letExp "reg_O" <=< segMap lvl (MkSolo ws_prod) $ \(MkSolo SubExp
i) -> do
    segment <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ws_data [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    offset <- letSubExp "offset" <=< toExp $ pe64 segment * pe64 segment_size
    pure [subExpRes offset]
  let rep' =
        IrregularRep
          { irregularS :: VName
irregularS = VName
arr_S,
            irregularF :: VName
irregularF = VName
arr_F,
            irregularO :: VName
irregularO = VName
arr_O,
            irregularD :: VName
irregularD = VName
arr_D,
            irregularK :: IrregularKind
irregularK = IrregularKind
Replicated
          }
  pure rep'

-- Replicates inner dimension for inputs.
onMapFreeVar ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  VName ->
  (VName, VName, VName) ->
  VName ->
  Maybe (FlattenM (VName, MapArray IrregularRep))
onMapFreeVar :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> VName
-> (VName, VName, VName)
-> VName
-> Maybe (FlattenM (VName, MapArray IrregularRep))
onMapFreeVar SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps VName
_ws (VName
_ws_F, VName
_ws_O, VName
ws_data) VName
v = do
  v_inp <- VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps
  pure $ do
    ws_prod <- arraySize 0 <$> lookupType ws_data
    fmap (v,) $ case v_inp of
      DistInputFree VName
v' Type
t ->
        --  I'm not totally sure if this will be better than previous approach
        (IrregularRep -> Type -> MapArray IrregularRep
forall t. t -> Type -> MapArray t
`MapOther` Type
t) (IrregularRep -> MapArray IrregularRep)
-> FlattenM IrregularRep -> FlattenM (MapArray IrregularRep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel -> Segments -> VName -> VName -> FlattenM IrregularRep
regularToReplicatedIrregularRep SegLevel
lvl Segments
segments VName
ws_data VName
v'
      DistInput ResTag
rt Type
t -> case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
        Irregular IrregularRep
rep -> do
          ~[new_S, offsets] <- Name -> Exp (Rep FlattenM) -> FlattenM [VName]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_rep_free_irreg")
            (Exp GPU -> FlattenM [VName])
-> ((Solo SubExp -> FlattenM Result) -> FlattenM (Exp GPU))
-> (Solo SubExp -> FlattenM Result)
-> FlattenM [VName]
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> Solo SubExp
-> (Solo SubExp -> 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 (SubExp -> Solo SubExp
forall a. a -> Solo a
MkSolo SubExp
ws_prod)
            ((Solo SubExp -> FlattenM Result) -> FlattenM [VName])
-> (Solo SubExp -> FlattenM Result) -> FlattenM [VName]
forall a b. (a -> b) -> a -> b
$ \(MkSolo SubExp
i) -> do
              segment <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ws_data [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
              s <- letSubExp "s" =<< eIndex (irregularS rep) [eSubExp segment]
              o <- letSubExp "o" =<< eIndex (irregularO rep) [eSubExp segment]
              pure $ subExpsRes [s, o]
          let rep' =
                IrregularRep
                  { irregularS :: VName
irregularS = VName
new_S,
                    irregularF :: VName
irregularF = IrregularRep -> VName
irregularF IrregularRep
rep,
                    irregularO :: VName
irregularO = VName
offsets,
                    irregularD :: VName
irregularD = IrregularRep -> VName
irregularD IrregularRep
rep,
                    irregularK :: IrregularKind
irregularK = IrregularKind
Replicated
                  }
          pure $ MapOther rep' t
        Regular VName
vs ->
          (IrregularRep -> Type -> MapArray IrregularRep
forall t. t -> Type -> MapArray t
`MapOther` Type
t) (IrregularRep -> MapArray IrregularRep)
-> FlattenM IrregularRep -> FlattenM (MapArray IrregularRep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel -> Segments -> VName -> VName -> FlattenM IrregularRep
regularToReplicatedIrregularRep SegLevel
lvl Segments
segments VName
ws_data VName
vs

onMapFreeVarMultiDim ::
  SegLevel ->
  Segments ->
  SubExp ->
  DistEnv ->
  DistInputs ->
  VName ->
  Maybe (FlattenM (VName, MapArray IrregularRep))
onMapFreeVarMultiDim :: SegLevel
-> Segments
-> SubExp
-> DistEnv
-> DistInputs
-> VName
-> Maybe (FlattenM (VName, MapArray IrregularRep))
onMapFreeVarMultiDim SegLevel
lvl Segments
segments SubExp
w DistEnv
env DistInputs
inps VName
v = do
  v_inp <- VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps
  pure $ fmap (v,) $ case v_inp of
    DistInputFree VName
v' Type
t -> do
      v_rep <- Segments -> ShapeBase SubExp -> VName -> FlattenM VName
replicateForDims Segments
segments (Segments -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) VName
v'
      pure $ MapArray v_rep t
    DistInput ResTag
rt Type
t -> case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
      Regular VName
v' -> do
        v_rep <- Segments -> ShapeBase SubExp -> VName -> FlattenM VName
replicateForDims Segments
segments (Segments -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) VName
v'
        pure $ MapArray v_rep t
      Irregular IrregularRep
rep -> do
        -- Can replicate as well
        old_nseg <- Int -> Type -> SubExp
forall o. Int -> TypeBase (ShapeBase SubExp) o -> SubExp
arraySize Int
0 (Type -> SubExp) -> FlattenM Type -> FlattenM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType (IrregularRep -> VName
irregularS IrregularRep
rep)
        new_nseg <- letSubExp "new_nseg" <=< toExp $ pe64 old_nseg * pe64 w
        ~[new_S, offsets] <- letTupExp (baseName v <> "_rep_free_irreg")
          <=< segMap lvl (MkSolo new_nseg)
          $ \(MkSolo SubExp
i) -> do
            old_seg <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"old_seg" (Exp GPU -> FlattenM SubExp)
-> (TPrimExp Int64 VName -> FlattenM (Exp GPU))
-> TPrimExp Int64 VName
-> FlattenM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
TPrimExp Int64 VName -> FlattenM (Exp GPU)
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp (TPrimExp Int64 VName -> FlattenM SubExp)
-> TPrimExp Int64 VName -> FlattenM SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
i TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`quot` SubExp -> TPrimExp Int64 VName
pe64 SubExp
w
            s <- letSubExp "s" =<< eIndex (irregularS rep) [eSubExp old_seg]
            o <- letSubExp "o" =<< eIndex (irregularO rep) [eSubExp old_seg]
            pure $ subExpsRes [s, o]
        let rep' =
              IrregularRep
                { irregularS :: VName
irregularS = VName
new_S,
                  irregularF :: VName
irregularF = IrregularRep -> VName
irregularF IrregularRep
rep,
                  irregularO :: VName
irregularO = VName
offsets,
                  irregularD :: VName
irregularD = IrregularRep -> VName
irregularD IrregularRep
rep,
                  irregularK :: IrregularKind
irregularK = IrregularKind
Replicated
                }
        pure $ MapOther rep' t

onMapInputArr ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  VName ->
  VName ->
  VName ->
  Param Type ->
  VName ->
  FlattenM (MapArray IrregularRep)
onMapInputArr :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> VName
-> VName
-> VName
-> Param Type
-> VName
-> FlattenM (MapArray IrregularRep)
onMapInputArr SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps VName
ws VName
ws_O VName
ws_data Param Type
p VName
arr = do
  ws_prod <- Int -> Type -> SubExp
forall o. Int -> TypeBase (ShapeBase SubExp) o -> SubExp
arraySize Int
0 (Type -> SubExp) -> FlattenM Type -> FlattenM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ws_data
  case lookup arr inps of
    Just DistInput
v_inp ->
      case DistInput
v_inp of
        DistInputFree VName
vs Type
t -> do
          let inner_shape :: ShapeBase SubExp
inner_shape = Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape (Type -> ShapeBase SubExp) -> Type -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p
          vs_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
vs
          v <-
            if isAcc vs_t
              then pure vs
              else
                letExp (baseName vs <> "_flat") . BasicOp . Reshape vs $
                  reshapeAll (arrayShape vs_t) (Shape [ws_prod] <> inner_shape)
          pure $ MapArray v (rowType t)
        DistInput ResTag
rt Type
t ->
          case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
            Irregular IrregularRep
rep -> do
              SegLevel
-> InnerMapMode
-> Segments
-> VName
-> VName
-> VName
-> Param Type
-> VName
-> IrregularRep
-> SubExp
-> FlattenM (MapArray IrregularRep)
onMapIrregularInputArr SegLevel
lvl InnerMapMode
SingleDim Segments
segments VName
ws VName
ws_O VName
ws_data Param Type
p VName
arr IrregularRep
rep SubExp
ws_prod
            Regular VName
vs -> do
              let inner_shape :: ShapeBase SubExp
inner_shape = Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape (Type -> ShapeBase SubExp) -> Type -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p
              vs_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
vs
              if isAcc vs_t
                then pure $ MapArray vs t
                else do
                  v <-
                    letExp (baseName arr <> "_reg_flat") . BasicOp . Reshape vs $
                      reshapeAll (arrayShape vs_t) (Shape [ws_prod] <> inner_shape)
                  pure $ MapArray v (stripArray 1 vs_t)
    Maybe DistInput
Nothing -> do
      arr_row_t <- Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType (Type -> Type) -> FlattenM Type -> FlattenM Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
      arr_rep <-
        letExp (baseName arr <> "_inp_rep") . BasicOp $
          Replicate (segmentsShape segments) (Var arr)
      arr_rep_t <- lookupType arr_rep
      v <-
        letExp (baseName arr <> "_inp_rep_flat") . BasicOp . Reshape arr_rep $
          reshapeAll (arrayShape arr_rep_t) (Shape [ws_prod] <> arrayShape arr_row_t)
      pure $ MapArray v arr_row_t

mapArraysToInputs ::
  [VName] ->
  [MapArray IrregularRep] ->
  FlattenM (DistEnv, DistInputs)
mapArraysToInputs :: [VName]
-> [MapArray IrregularRep] -> FlattenM (DistEnv, DistInputs)
mapArraysToInputs [VName]
param_names [MapArray IrregularRep]
arrs = do
  ((_, env), inputs) <-
    ((Int, DistEnv)
 -> (VName, MapArray IrregularRep)
 -> FlattenM ((Int, DistEnv), (VName, DistInput)))
-> (Int, DistEnv)
-> [(VName, MapArray IrregularRep)]
-> FlattenM ((Int, DistEnv), DistInputs)
forall (m :: * -> *) (t :: * -> *) acc x y.
(Monad m, Traversable t) =>
(acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
mapAccumLM (Int, DistEnv)
-> (VName, MapArray IrregularRep)
-> FlattenM ((Int, DistEnv), (VName, DistInput))
forall {a}.
(Int, DistEnv)
-> (a, MapArray IrregularRep)
-> FlattenM ((Int, DistEnv), (a, DistInput))
onInput (Int
0, DistEnv
forall a. Monoid a => a
mempty) ([(VName, MapArray IrregularRep)]
 -> FlattenM ((Int, DistEnv), DistInputs))
-> [(VName, MapArray IrregularRep)]
-> FlattenM ((Int, DistEnv), DistInputs)
forall a b. (a -> b) -> a -> b
$ [VName]
-> [MapArray IrregularRep] -> [(VName, MapArray IrregularRep)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
param_names [MapArray IrregularRep]
arrs
  pure (env, inputs)
  where
    onInput :: (Int, DistEnv)
-> (a, MapArray IrregularRep)
-> FlattenM ((Int, DistEnv), (a, DistInput))
onInput (Int
tag, DistEnv
env) (a
p, MapArray VName
arr Type
t) =
      ((Int, DistEnv), (a, DistInput))
-> FlattenM ((Int, DistEnv), (a, DistInput))
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int
tag, DistEnv
env), (a
p, VName -> Type -> DistInput
DistInputFree VName
arr Type
t))
    onInput (Int
tag, DistEnv
env) (a
p, MapOther IrregularRep
rep Type
t) = do
      let rt :: ResTag
rt = Int -> ResTag
ResTag Int
tag
      env' <- ResTag -> ResRep -> DistEnv -> FlattenM DistEnv
insertRepM ResTag
rt (IrregularRep -> ResRep
Irregular IrregularRep
rep) DistEnv
env
      pure
        ( (tag + 1, env'),
          (p, DistInput rt t)
        )

transformUniformRedomap ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  SubExp ->
  [VName] ->
  [Reduce SOACS] ->
  Lambda SOACS ->
  FlattenM [VName]
transformUniformRedomap :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> SubExp
-> [VName]
-> [Reduce SOACS]
-> Lambda SOACS
-> FlattenM [VName]
transformUniformRedomap SegLevel
lvl [] DistEnv
_env DistInputs
_inps SubExp
w [VName]
arrs [Reduce SOACS]
reds Lambda SOACS
map_lam = do
  -- Top-level (no enclosing segments): the arrays and any free variables are
  -- ordinary top-level values in scope, so this is an ordinary non-segmented
  -- reduce over the map width. We emit it as such ('genNonSegRed'); a segmented
  -- reduce over a single implicit segment would be equivalent, but downstream
  -- passes (e.g. migration, coalescing) handle the non-segmented form better.
  let sing_red :: Reduce SOACS
sing_red = [Reduce SOACS] -> Reduce SOACS
forall rep. Buildable rep => [Reduce rep] -> Reduce rep
singleReduce [Reduce SOACS]
reds
  (red_lam, nes', shape) <- Lambda SOACS
-> Segments -> FlattenM (Lambda SOACS, Segments, ShapeBase SubExp)
forall (m :: * -> *).
MonadBuilder m =>
Lambda SOACS
-> Segments -> m (Lambda SOACS, Segments, ShapeBase SubExp)
determineReduceOp (Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda Reduce SOACS
sing_red) (Reduce SOACS -> Segments
forall rep. Reduce rep -> Segments
redNeutral Reduce SOACS
sing_red)
  let comm
        | Lambda SOACS -> Bool
forall rep. Lambda rep -> Bool
commutativeLambda Lambda SOACS
red_lam = Commutativity
Commutative
        | Bool
otherwise = Reduce SOACS -> Commutativity
forall rep. Reduce rep -> Commutativity
redComm Reduce SOACS
sing_red
      sing_red_gpu = Commutativity -> GLambda GPU Type -> Segments -> Reduce GPU
forall rep. Commutativity -> Lambda rep -> Segments -> Reduce rep
Reduce Commutativity
comm (Lambda SOACS -> GLambda GPU Type
soacsLambdaToGPU Lambda SOACS
red_lam) Segments
nes'
  genNonSegRed lvl "topLevelSegRed" [w] sing_red_gpu shape (soacsLambdaToGPU map_lam) arrs
transformUniformRedomap SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps SubExp
w [VName]
arrs [Reduce SOACS]
reds Lambda SOACS
map_lam = do
  let free :: Names
free = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
map_lam
      new_segment :: Segments
new_segment = Segments
segments Segments -> Segments -> Segments
forall a. Semigroup a => a -> a -> a
<> SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w
      shape :: ShapeBase SubExp
shape = ShapeBase SubExp
forall a. Monoid a => a
mempty
  reds_gpu <- [Reduce SOACS]
-> (Reduce SOACS -> FlattenM (Reduce GPU)) -> FlattenM [Reduce GPU]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Reduce SOACS]
reds ((Reduce SOACS -> FlattenM (Reduce GPU)) -> FlattenM [Reduce GPU])
-> (Reduce SOACS -> FlattenM (Reduce GPU)) -> FlattenM [Reduce GPU]
forall a b. (a -> b) -> a -> b
$ \Reduce SOACS
red -> do
    nes <- (SubExp -> FlattenM SubExp) -> Segments -> FlattenM Segments
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 (Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM SubExp
readNeutral Segments
segments DistEnv
env DistInputs
inps) (Reduce SOACS -> Segments
forall rep. Reduce rep -> Segments
redNeutral Reduce SOACS
red)
    let red_lam = Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda Reduce SOACS
red
        comm
          | Lambda SOACS -> Bool
forall rep. Lambda rep -> Bool
commutativeLambda Lambda SOACS
red_lam = Commutativity
Commutative
          | Bool
otherwise = Reduce SOACS -> Commutativity
forall rep. Reduce rep -> Commutativity
redComm Reduce SOACS
red
    pure $ Reduce comm (soacsLambdaToGPU red_lam) nes
  free_and_sizes <- freeWithTypeDeps inps free
  (free_replicated, replicated) <-
    fmap unzip . sequence $
      mapMaybe
        (onMapFreeVarMultiDim lvl segments w env inps)
        free_and_sizes
  arrs' <-
    zipWithM
      ( \Param Type
p VName
arr ->
          SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> ShapeBase SubExp
-> SubExp
-> FlattenM VName
liftSubExpRegular
            SegLevel
lvl
            Segments
segments
            DistInputs
inps
            DistEnv
env
            (Segments -> ShapeBase SubExp
segmentsShape Segments
new_segment ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape (Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p))
            (VName -> SubExp
Var VName
arr)
      )
      (lambdaParams map_lam)
      arrs

  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated
  let readFree Segments
is = Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs Segments
new_segment DistEnv
free_env Segments
is DistInputs
free_inputs
  genUniformSegRed lvl "uniformSegRed" new_segment reds_gpu shape (soacsLambdaToGPU map_lam) arrs' readFree

doUniformSegMaposcanomap ::
  SegLevel ->
  [Scan SOACS] ->
  [VName] ->
  Lambda SOACS ->
  Lambda SOACS ->
  Segments ->
  Segments ->
  DistInputs ->
  DistEnv ->
  ([SubExp] -> FlattenM ()) ->
  FlattenM [VName]
doUniformSegMaposcanomap :: SegLevel
-> [Scan SOACS]
-> [VName]
-> Lambda SOACS
-> Lambda SOACS
-> Segments
-> Segments
-> DistInputs
-> DistEnv
-> (Segments -> FlattenM ())
-> FlattenM [VName]
doUniformSegMaposcanomap SegLevel
lvl [Scan SOACS]
scans [VName]
arrs Lambda SOACS
post_lam Lambda SOACS
map_lam Segments
old_segments Segments
new_segment DistInputs
inps DistEnv
env Segments -> FlattenM ()
readFree = do
  let scan :: Scan SOACS
scan = [Scan SOACS] -> Scan SOACS
forall rep. Buildable rep => [Scan rep] -> Scan rep
singleScan [Scan SOACS]
scans
  nes <- (SubExp -> FlattenM SubExp) -> Segments -> FlattenM Segments
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 (Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM SubExp
readNeutral Segments
old_segments DistEnv
env DistInputs
inps) (Scan SOACS -> Segments
forall rep. Scan rep -> Segments
scanNeutral Scan SOACS
scan)
  (scan_lam, nes', shape) <- determineReduceOp (scanLambda scan) nes
  genUniformSegScanomapWithPost
    lvl
    new_segment
    "uniformmaposcanomap"
    (soacsLambdaToGPU scan_lam)
    shape
    nes'
    (soacsLambdaToGPU post_lam)
    (soacsLambdaToGPU map_lam)
    arrs
    readFree

transformUniformMaposcanomap ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  SubExp ->
  [VName] ->
  [Scan SOACS] ->
  Lambda SOACS ->
  Lambda SOACS ->
  FlattenM [VName]
transformUniformMaposcanomap :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> SubExp
-> [VName]
-> [Scan SOACS]
-> Lambda SOACS
-> Lambda SOACS
-> FlattenM [VName]
transformUniformMaposcanomap SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps SubExp
w [VName]
arrs [Scan SOACS]
scans Lambda SOACS
post_lam Lambda SOACS
map_lam = do
  let free :: Names
free = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
map_lam Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
post_lam
      new_segment :: Segments
new_segment = Segments
segments Segments -> Segments -> Segments
forall a. Semigroup a => a -> a -> a
<> SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w
  free_and_sizes <- DistInputs -> Names -> FlattenM [VName]
freeWithTypeDeps DistInputs
inps Names
free
  (free_replicated, replicated) <-
    fmap unzip . sequence $
      mapMaybe
        (onMapFreeVarMultiDim lvl segments w env inps)
        free_and_sizes
  arrs' <-
    zipWithM
      ( \Param Type
p VName
arr ->
          SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> ShapeBase SubExp
-> SubExp
-> FlattenM VName
liftSubExpRegular
            SegLevel
lvl
            Segments
segments
            DistInputs
inps
            DistEnv
env
            (Segments -> ShapeBase SubExp
segmentsShape Segments
new_segment ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape (Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p))
            (VName -> SubExp
Var VName
arr)
      )
      (lambdaParams map_lam)
      arrs
  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated
  let readFree Segments
is = Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs Segments
new_segment DistEnv
free_env Segments
is DistInputs
free_inputs
  doUniformSegMaposcanomap lvl scans arrs' post_lam map_lam segments new_segment inps env readFree

doSegMaposcanomap ::
  SegLevel ->
  [Scan SOACS] ->
  VName ->
  [VName] ->
  Lambda SOACS ->
  Lambda SOACS ->
  Segments ->
  DistInputs ->
  DistEnv ->
  ([SubExp] -> FlattenM ()) ->
  FlattenM [VName]
doSegMaposcanomap :: SegLevel
-> [Scan SOACS]
-> VName
-> [VName]
-> Lambda SOACS
-> Lambda SOACS
-> Segments
-> DistInputs
-> DistEnv
-> (Segments -> FlattenM ())
-> FlattenM [VName]
doSegMaposcanomap SegLevel
lvl [Scan SOACS]
scans VName
flags [VName]
elems Lambda SOACS
post_lam Lambda SOACS
map_lam Segments
segments DistInputs
inps DistEnv
env Segments -> FlattenM ()
readFree = do
  let scan :: Scan SOACS
scan = [Scan SOACS] -> Scan SOACS
forall rep. Buildable rep => [Scan rep] -> Scan rep
singleScan [Scan SOACS]
scans
  let nes :: Segments
nes = Scan SOACS -> Segments
forall rep. Scan rep -> Segments
scanNeutral Scan SOACS
scan
  nes' <- (SubExp -> FlattenM SubExp) -> Segments -> FlattenM Segments
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 (Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM SubExp
readNeutral Segments
segments DistEnv
env DistInputs
inps) Segments
nes
  genSegScanomapWithPost
    lvl
    "maposcanomap"
    (soacsLambdaToGPU $ scanLambda scan)
    nes'
    flags
    (soacsLambdaToGPU post_lam)
    (soacsLambdaToGPU map_lam)
    elems
    readFree

-- Hacky fix to get result representations in the same order as the pattern
resRepsInPatOrder :: Pat Type -> [(VName, ResRep)] -> [ResRep]
resRepsInPatOrder :: Pat Type -> [(VName, ResRep)] -> [ResRep]
resRepsInPatOrder Pat Type
pat [(VName, ResRep)]
reps =
  let rep_map :: Map VName ResRep
rep_map = [(VName, ResRep)] -> Map VName ResRep
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, ResRep)]
reps
      lookupRes :: VName -> ResRep
lookupRes VName
v =
        case VName -> Map VName ResRep -> Maybe ResRep
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName ResRep
rep_map of
          Just ResRep
rep -> ResRep
rep
          Maybe ResRep
Nothing ->
            String -> ResRep
forall a. HasCallStack => String -> a
error (String -> ResRep) -> String -> ResRep
forall a b. (a -> b) -> a -> b
$
              String
"resRepsInPatOrder: missing result for "
                String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
v
   in (VName -> ResRep) -> [VName] -> [ResRep]
forall a b. (a -> b) -> [a] -> [b]
map VName -> ResRep
lookupRes (Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pat)

segOpInputRep ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  VName ->
  FlattenM ResRep
segOpInputRep :: SegLevel
-> Segments -> DistEnv -> DistInputs -> VName -> FlattenM ResRep
segOpInputRep SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps VName
arr =
  case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
arr DistInputs
inps of
    Just (DistInput ResTag
rt Type
_) ->
      ResRep -> FlattenM ResRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResRep -> FlattenM ResRep) -> ResRep -> FlattenM ResRep
forall a b. (a -> b) -> a -> b
$ ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env
    Just (DistInputFree VName
arr' Type
_) ->
      ResRep -> FlattenM ResRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResRep -> FlattenM ResRep) -> ResRep -> FlattenM ResRep
forall a b. (a -> b) -> a -> b
$ VName -> ResRep
Regular VName
arr'
    Maybe DistInput
Nothing ->
      IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep)
-> FlattenM IrregularRep -> FlattenM ResRep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps VName
arr

-- Basically we need to make our arrays ready for our segscan/segred.
-- Regular arrays are flattened only across the outer segment dimensions and
-- the SOAC width; any row shape expected by the consumer is preserved.
-- we need to check the dense/replicated status of the input.
-- if all of scan inputs are replicated we are fine.
-- otherwise, we need to make the replicated inputs dense.
-- for regulars we can just use the segment descriptor and this should be also the same descriptor for dense irregulars.
prepareSegOpInputs ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  SubExp ->
  [ResRep] ->
  [VName] ->
  Bool ->
  FlattenM (VName, VName, VName, [VName], IrregularKind)
prepareSegOpInputs :: SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> SubExp
-> [ResRep]
-> [VName]
-> Bool
-> FlattenM (VName, VName, VName, [VName], IrregularKind)
prepareSegOpInputs SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps SubExp
w [ResRep]
reps [VName]
names Bool
hasNoFreeVariant
  | (ResRep -> Bool) -> [ResRep] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ResRep -> Bool
isRegular [ResRep]
reps = do
      ws <- SegLevel
-> Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps SubExp
w
      (ws_F, ws_O, ws_data) <- doRepIota lvl ws
      m <- arraySize 0 <$> lookupType ws_data
      names' <- mapM (flattenRegularRep m) reps
      pure (ws_F, ws_O, ws, names', Dense)
  | (ResRep -> Bool) -> [ResRep] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ResRep -> Bool
isReplicatedIrregular [ResRep]
reps Bool -> Bool -> Bool
&& Bool
hasNoFreeVariant = do
      -- We use the descriptor of the first rep for all inputs, which assumes
      -- that all the replicated inputs have the same offsets into their
      -- respective data arrays. This holds because same-width views produced
      -- by onMapFreeVar inherit the offsets of their underlying arrays, and
      -- those are compact per-segment arrays of the SOAC width (slices are
      -- materialised before they can become inputs here).
      let Irregular IrregularRep
rep0 = [ResRep] -> ResRep
forall a. HasCallStack => [a] -> a
head [ResRep]
reps
      (VName, VName, VName, [VName], IrregularKind)
-> FlattenM (VName, VName, VName, [VName], IrregularKind)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IrregularRep -> VName
irregularF IrregularRep
rep0, IrregularRep -> VName
irregularO IrregularRep
rep0, IrregularRep -> VName
irregularS IrregularRep
rep0, (ResRep -> VName) -> [ResRep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ResRep -> VName
getData [ResRep]
reps, IrregularKind
Replicated)
  | Bool
otherwise = do
      -- The segment descriptor must count SOAC *elements* per segment, but an
      -- 'IrregularRep' stores primitive data whose structure arrays count
      -- _scalars_ - @c@ per element for a non-scalar element type @b=[c]t@. We
      -- take a dense input rep as the descriptor and convert its scalar-unit
      -- structure to element units: sizes and offsets divided by that input's
      -- @c@, flags subsampled with stride @c@ (all no-ops when @c=1@). Each
      -- input's flat data is reshaped into @[m]b@ rows.
      row_types <- (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 ((Type -> Type) -> FlattenM Type -> FlattenM Type
forall a b. (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType (FlattenM Type -> FlattenM Type)
-> (VName -> FlattenM Type) -> VName -> FlattenM Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps) [VName]
names
      (desc, desc_c) <- descriptor $ zip reps row_types
      ws_S <- scaleSizesDown lvl desc_c (irregularS desc)
      ws_O <- scaleSizesDown lvl desc_c (irregularO desc)
      ws_F <- subsampleFlags lvl desc_c (irregularF desc)
      m <- arraySize 0 <$> lookupType ws_F
      names' <- sequence $ zipWith3 (toRows m) reps names row_types
      pure (ws_F, ws_O, ws_S, names', Dense)
  where
    isRegular :: ResRep -> Bool
isRegular (Regular VName
_) = Bool
True
    isRegular ResRep
_ = Bool
False

    isReplicatedIrregular :: ResRep -> Bool
isReplicatedIrregular (Irregular IrregularRep
rep) = IrregularRep -> IrregularKind
irregularK IrregularRep
rep IrregularKind -> IrregularKind -> Bool
forall a. Eq a => a -> a -> Bool
== IrregularKind
Replicated
    isReplicatedIrregular ResRep
_ = Bool
False

    flattenRegularRep :: SubExp -> ResRep -> FlattenM VName
flattenRegularRep SubExp
m (Regular VName
v) =
      Segments -> SubExp -> VName -> FlattenM VName
flattenRegularToRows Segments
segments SubExp
m VName
v
    flattenRegularRep SubExp
_ ResRep
_ =
      String -> FlattenM VName
forall a. HasCallStack => String -> a
error String
"prepareSegOpInputs: impossible irregular regular input"
    getData :: ResRep -> VName
getData (Irregular IrregularRep
rep) = IrregularRep -> VName
irregularD IrregularRep
rep
    getData ResRep
_ = String -> VName
forall a. HasCallStack => String -> a
error String
"prepareSegOpInputs: impossible"

    -- A dense input rep to use as the descriptor, and its element type's inner
    -- size @c@. Any dense irregular input serves, since all inputs share the
    -- SOAC's segmentation.
    descriptor :: [(ResRep, TypeBase (ShapeBase SubExp) m)]
-> FlattenM (IrregularRep, TPrimExp Int64 VName)
descriptor [(ResRep, TypeBase (ShapeBase SubExp) m)]
rs =
      case [(IrregularRep
rep, TypeBase (ShapeBase SubExp) m
t) | (Irregular IrregularRep
rep, TypeBase (ShapeBase SubExp) m
t) <- [(ResRep, TypeBase (ShapeBase SubExp) m)]
rs, IrregularRep -> IrregularKind
irregularK IrregularRep
rep IrregularKind -> IrregularKind -> Bool
forall a. Eq a => a -> a -> Bool
== IrregularKind
Dense] of
        (IrregularRep
rep, TypeBase (ShapeBase SubExp) m
t) : [(IrregularRep, TypeBase (ShapeBase SubExp) m)]
_ -> (IrregularRep, TPrimExp Int64 VName)
-> FlattenM (IrregularRep, TPrimExp Int64 VName)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IrregularRep
rep, TypeBase (ShapeBase SubExp) m -> TPrimExp Int64 VName
forall {m}. TypeBase (ShapeBase SubExp) m -> TPrimExp Int64 VName
innerSize TypeBase (ShapeBase SubExp) m
t)
        [] ->
          case [(IrregularRep
rep, TypeBase (ShapeBase SubExp) m
t) | (Irregular IrregularRep
rep, TypeBase (ShapeBase SubExp) m
t) <- [(ResRep, TypeBase (ShapeBase SubExp) m)]
rs] of
            (IrregularRep
rep, TypeBase (ShapeBase SubExp) m
t) : [(IrregularRep, TypeBase (ShapeBase SubExp) m)]
_ -> do
              rep' <- SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl Name
"segop_desc" IrregularRep
rep
              pure (rep', innerSize t)
            [] -> String -> FlattenM (IrregularRep, TPrimExp Int64 VName)
forall a. HasCallStack => String -> a
error String
"prepareSegOpInputs: impossible"
    innerSize :: TypeBase (ShapeBase SubExp) m -> TPrimExp Int64 VName
innerSize TypeBase (ShapeBase SubExp) m
t = [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ (SubExp -> TPrimExp Int64 VName)
-> Segments -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 (Segments -> [TPrimExp Int64 VName])
-> Segments -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase SubExp) m -> Segments
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims TypeBase (ShapeBase SubExp) m
t

    -- Reshape an input's data into @m@ rows of the (possibly non-scalar) element
    -- type. For an irregular input the flat scalar data is grouped into rows;
    -- for a regular input the enclosing dimensions are collapsed.
    toRows :: SubExp
-> ResRep
-> VName
-> TypeBase (ShapeBase SubExp) o
-> FlattenM VName
toRows SubExp
m (Regular VName
v') VName
_ TypeBase (ShapeBase SubExp) o
_ =
      Segments -> SubExp -> VName -> FlattenM VName
flattenRegularToRows Segments
segments SubExp
m VName
v'
    toRows SubExp
m (Irregular IrregularRep
ir) VName
v TypeBase (ShapeBase SubExp) o
row_t = do
      d <- IrregularRep -> VName
irregularD (IrregularRep -> VName) -> FlattenM IrregularRep -> FlattenM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_dense") IrregularRep
ir
      d_t <- lookupType d
      letExp (baseName v <> "_rows") . BasicOp $
        Reshape d $
          reshapeAll (arrayShape d_t) (Shape [m] <> arrayShape row_t)

flattenRegularToRows :: Segments -> SubExp -> VName -> FlattenM VName
flattenRegularToRows :: Segments -> SubExp -> VName -> FlattenM VName
flattenRegularToRows Segments
segments SubExp
m VName
v = do
  v_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
  if isAcc v_t
    then pure v
    else do
      when (arrayRank v_t < segmentsRank segments + 1) $
        error "prepareSegOpInputs: regular input rank too small"
      let row_shape = Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape (Type -> ShapeBase SubExp) -> Type -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Type -> Type
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray (Segments -> Int
segmentsRank Segments
segments Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Type
v_t
      letExp (baseName v <> "_flat") . BasicOp $
        Reshape v $
          reshapeAll (arrayShape v_t) (Shape [m] <> row_shape)

-- | Construct a body and immediately rename it.
renamedBody :: FlattenM [VName] -> FlattenM (Body GPU)
renamedBody :: FlattenM [VName] -> FlattenM (Body GPU)
renamedBody = Body GPU -> FlattenM (Body GPU)
forall rep res (m :: * -> *).
(Renameable rep, Rename res, MonadFreshNames m) =>
GBody rep res -> m (GBody rep res)
renameBody (Body GPU -> FlattenM (Body GPU))
-> (FlattenM [VName] -> FlattenM (Body GPU))
-> FlattenM [VName]
-> FlattenM (Body GPU)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< FlattenM Result -> FlattenM (GBody (Rep FlattenM) SubExpRes)
FlattenM Result -> FlattenM (Body GPU)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (FlattenM Result -> FlattenM (Body GPU))
-> (FlattenM [VName] -> FlattenM Result)
-> FlattenM [VName]
-> FlattenM (Body GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([VName] -> Result) -> FlattenM [VName] -> FlattenM Result
forall a b. (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [VName] -> Result
varsRes

regularRepVars :: [ResRep] -> [VName]
regularRepVars :: [ResRep] -> [VName]
regularRepVars =
  (ResRep -> VName) -> [ResRep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ResRep -> VName
onRep
  where
    onRep :: ResRep -> VName
onRep (Regular VName
v) = VName
v
    onRep Irregular {} = String -> VName
forall a. HasCallStack => String -> a
error String
"regularRepVars: expected regular result"

versionScanRed ::
  FlattenOps ->
  Name ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  StmAux () ->
  SubExp ->
  Body SOACS ->
  FlattenM [VName] ->
  FlattenM DistEnv
versionScanRed :: FlattenOps
-> Name
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> SubExp
-> Body SOACS
-> FlattenM [VName]
-> FlattenM DistEnv
versionScanRed FlattenOps
ops Name
desc Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux SubExp
w Body SOACS
factored_body FlattenM [VName]
outer_only = do
  let result_ts :: [Type]
result_ts =
        [ Type
t Type -> ShapeBase SubExp -> Type
`arrayOfShape` Segments -> ShapeBase SubExp
segmentsShape Segments
segments
        | DistResult ResTag
_ (DistType Segments
_ Rank
_ Type
t) VName
_ <- [DistResult]
res
        ]
  let fullFlatten :: FlattenM (Body GPU)
fullFlatten = case Segments
segments of
        -- Top-level (no enclosing segments): flatten the factored body's
        -- statements as ordinary top-level statements. Unlike distributing them
        -- over segments, this copes with array-valued operators and nested
        -- SOACs whose temporaries would otherwise escape the segmented
        -- machinery's scope.
        [] ->
          Body GPU -> FlattenM (Body GPU)
forall rep res (m :: * -> *).
(Renameable rep, Rename res, MonadFreshNames m) =>
GBody rep res -> m (GBody rep res)
renameBody (Body GPU -> FlattenM (Body GPU))
-> (FlattenM Result -> FlattenM (Body GPU))
-> FlattenM Result
-> FlattenM (Body GPU)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< FlattenM Result -> FlattenM (GBody (Rep FlattenM) SubExpRes)
FlattenM Result -> FlattenM (Body GPU)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (FlattenM Result -> FlattenM (Body GPU))
-> FlattenM Result -> FlattenM (Body GPU)
forall a b. (a -> b) -> a -> b
$ do
            (Stm SOACS -> FlattenM ()) -> Seq (Stm SOACS) -> FlattenM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (FlattenOps -> Stm SOACS -> FlattenM ()
flattenTopLevelStm FlattenOps
ops) (Seq (Stm SOACS) -> FlattenM ()) -> Seq (Stm SOACS) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Seq (Stm SOACS)
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
factored_body
            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
$ Body SOACS -> Result
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
factored_body
        Segments
_ ->
          FlattenM [VName] -> FlattenM (Body GPU)
renamedBody (FlattenM [VName] -> FlattenM (Body GPU))
-> FlattenM [VName] -> FlattenM (Body GPU)
forall a b. (a -> b) -> a -> b
$ [ResRep] -> [VName]
regularRepVars ([ResRep] -> [VName]) -> FlattenM [ResRep] -> FlattenM [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FlattenOps
-> Segments
-> Name
-> DistEnv
-> DistInputs
-> [DistResult]
-> Body SOACS
-> FlattenM [ResRep]
distributeAndFlattenBody FlattenOps
ops Segments
segments Name
"versionScanRed_full_body" DistEnv
env DistInputs
inps [DistResult]
res Body SOACS
factored_body

  match_res <-
    Certs -> FlattenM [VName] -> FlattenM [VName]
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying (DistInputs -> StmAux () -> DistEnv -> Certs
forall a. DistInputs -> StmAux a -> DistEnv -> Certs
distCerts DistInputs
inps StmAux ()
aux DistEnv
env) (FlattenM [VName] -> FlattenM [VName])
-> FlattenM [VName] -> FlattenM [VName]
forall a b. (a -> b) -> a -> b
$
      Name
-> [Type]
-> Attrs
-> Bool
-> Bool
-> Segments
-> FlattenM (Body GPU)
-> FlattenM (Body GPU)
-> FlattenM [VName]
scanRedAlternatives
        Name
desc
        [Type]
result_ts
        (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
aux)
        (FunHasParallelism -> Body SOACS -> Bool
isParallelFunInside (FlattenOps -> FunHasParallelism
flattenFunHasParallelism FlattenOps
ops) Body SOACS
factored_body)
        (SegLevel -> Bool
allowVersioning (FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops))
        (Segments
segments Segments -> Segments -> Segments
forall a. Semigroup a => a -> a -> a
<> SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w)
        FlattenM (Body GPU)
fullFlatten
        (FlattenM [VName] -> FlattenM (Body GPU)
renamedBody FlattenM [VName]
outer_only)
  pure $ insertRegulars (map distResTag res) match_res env

insertSegOpMapResults ::
  Segments ->
  VName ->
  VName ->
  VName ->
  IrregularKind ->
  [(DistResult, VName)] ->
  DistEnv ->
  FlattenM DistEnv
insertSegOpMapResults :: Segments
-> VName
-> VName
-> VName
-> IrregularKind
-> [(DistResult, VName)]
-> DistEnv
-> FlattenM DistEnv
insertSegOpMapResults Segments
segments VName
segs VName
flags VName
offsets IrregularKind
kind [(DistResult, VName)]
bnds DistEnv
env0 =
  (DistEnv -> (DistResult, VName) -> FlattenM DistEnv)
-> DistEnv -> [(DistResult, VName)] -> FlattenM DistEnv
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM DistEnv -> (DistResult, VName) -> FlattenM DistEnv
insert DistEnv
env0 [(DistResult, VName)]
bnds
  where
    insert :: DistEnv -> (DistResult, VName) -> FlattenM DistEnv
insert DistEnv
env (DistResult
dist_res, VName
v)
      | DistResult -> Bool
isRegularDistResult DistResult
dist_res = do
          let DistType Segments
_ Rank
_ Type
t = DistResult -> DistType
distResType DistResult
dist_res
          if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
            then DistEnv -> FlattenM DistEnv
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DistEnv -> FlattenM DistEnv) -> DistEnv -> FlattenM DistEnv
forall a b. (a -> b) -> a -> b
$ [ResTag] -> [VName] -> DistEnv -> DistEnv
insertRegulars [DistResult -> ResTag
distResTag DistResult
dist_res] [VName
v] DistEnv
env
            else do
              let expected_shape :: ShapeBase SubExp
expected_shape = Segments -> ShapeBase SubExp
segmentsShape Segments
segments ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t
              v_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
              v' <-
                letExp (baseName v <> "_reshaped") . BasicOp $
                  Reshape v $
                    reshapeAll (arrayShape v_t) expected_shape
              pure $ insertRegulars [distResTag dist_res] [v'] env
      | Bool
otherwise =
          VName
-> VName
-> VName
-> ResTag
-> VName
-> IrregularKind
-> DistEnv
-> FlattenM DistEnv
insertIrregularM VName
segs VName
flags VName
offsets (DistResult -> ResTag
distResTag DistResult
dist_res) VName
v IrregularKind
kind DistEnv
env

distResCerts :: DistEnv -> [DistInput] -> Certs
distResCerts :: DistEnv -> [DistInput] -> Certs
distResCerts DistEnv
env = [VName] -> Certs
Certs ([VName] -> Certs)
-> ([DistInput] -> [VName]) -> [DistInput] -> Certs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DistInput -> VName) -> [DistInput] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map DistInput -> VName
f
  where
    f :: DistInput -> VName
f (DistInputFree VName
v Type
_) = VName
v
    f (DistInput ResTag
rt Type
_) = case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
      Regular VName
v -> VName
v
      Irregular IrregularRep
r -> IrregularRep -> VName
irregularD IrregularRep
r

reshapeAndBind :: VName -> VName -> Shape -> FlattenM ()
reshapeAndBind :: VName -> VName -> ShapeBase SubExp -> FlattenM ()
reshapeAndBind VName
v VName
src ShapeBase SubExp
shape = do
  v_copy <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName VName
v) (Exp GPU -> FlattenM VName)
-> (BasicOp -> Exp GPU) -> BasicOp -> FlattenM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp GPU
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> FlattenM VName) -> BasicOp -> FlattenM VName
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
src
  v_copy_shape <- arrayShape <$> lookupType v_copy
  letBindNames [v] $ BasicOp $ Reshape v_copy $ reshapeAll v_copy_shape shape

mapResultRep :: SegLevel -> InnerMapMode -> (VName, VName, VName) -> VName -> FlattenM ResRep
mapResultRep :: SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> VName
-> FlattenM ResRep
mapResultRep SegLevel
_ InnerMapMode
MultiDim (VName, VName, VName)
_ VName
v = ResRep -> FlattenM ResRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResRep -> FlattenM ResRep) -> ResRep -> FlattenM ResRep
forall a b. (a -> b) -> a -> b
$ VName -> ResRep
Regular VName
v
mapResultRep SegLevel
lvl InnerMapMode
SingleDim (VName
ws, VName
ws_F, VName
ws_O) VName
v =
  -- Forcing the irregular rep to be 1D because in some places that is my assumption
  -- and also this will make the metadata consistent.
  IrregularRep -> ResRep
Irregular
    (IrregularRep -> ResRep)
-> FlattenM IrregularRep -> FlattenM ResRep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel -> IrregularRep -> FlattenM IrregularRep
flattenIrregularRep
      SegLevel
lvl
      IrregularRep
        { irregularS :: VName
irregularS = VName
ws,
          irregularF :: VName
irregularF = VName
ws_F,
          irregularO :: VName
irregularO = VName
ws_O,
          irregularD :: VName
irregularD = VName
v,
          irregularK :: IrregularKind
irregularK = IrregularKind
Dense
        }

transformDistributed ::
  FlattenOps ->
  M.Map ResTag IrregularRep ->
  Segments ->
  Distributed ->
  FlattenM ()
transformDistributed :: FlattenOps
-> Map ResTag IrregularRep
-> Segments
-> Distributed
-> FlattenM ()
transformDistributed FlattenOps
ops Map ResTag IrregularRep
irregs Segments
segments Distributed
dist = do
  let Distributed DistStms
dstms (DistResults ResMap
resmap [DistRep]
reps) = Distributed
dist
  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_initial DistStms
dstms
  forM_ (M.toList resmap) $ \(ResTag
rt, [([DistInput], VName, Type)]
binds) ->
    [([DistInput], VName, Type)]
-> (([DistInput], VName, Type) -> FlattenM ()) -> FlattenM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [([DistInput], VName, Type)]
binds ((([DistInput], VName, Type) -> FlattenM ()) -> FlattenM ())
-> (([DistInput], VName, Type) -> FlattenM ()) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ \([DistInput]
cs_inps, VName
v, Type
v_t) ->
      Certs -> FlattenM () -> FlattenM ()
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying (DistEnv -> [DistInput] -> Certs
distResCerts DistEnv
env [DistInput]
cs_inps) (FlattenM () -> FlattenM ()) -> FlattenM () -> FlattenM ()
forall a b. (a -> b) -> a -> b
$
        case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
          Regular 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'
          Irregular IrregularRep
irreg -> do
            -- It might have an irregular representation, but we know
            -- that it is actually regular because it is a result.
            irreg' <- SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular (FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops) (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_dist_res") IrregularRep
irreg
            reshapeAndBind v (irregularD irreg') (segmentsShape segments <> arrayShape v_t)
  forM_ reps $ \(VName
v, Either SubExp DistInput
r) ->
    case Either SubExp DistInput
r of
      Left SubExp
se ->
        [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
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate (Segments -> ShapeBase SubExp
segmentsShape Segments
segments) SubExp
se
      Right (DistInputFree VName
arr Type
_) ->
        [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
arr
      -- This can happen. ask Troels
      Right (DistInput ResTag
rt Type
t) ->
        case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
          Regular 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'
          Irregular IrregularRep
irreg -> do
            irreg' <- SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular (FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops) (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_dist_rep") IrregularRep
irreg
            reshapeAndBind v (irregularD irreg') (segmentsShape segments <> arrayShape t)
  where
    env_initial :: DistEnv
env_initial = DistEnv {distResMap :: Map ResTag ResRep
distResMap = (IrregularRep -> ResRep)
-> Map ResTag IrregularRep -> Map ResTag ResRep
forall a b k. (a -> b) -> Map k a -> Map k b
M.map IrregularRep -> ResRep
Irregular Map ResTag IrregularRep
irregs}

onMapIrregularInputArr ::
  SegLevel ->
  InnerMapMode ->
  Segments ->
  VName ->
  VName ->
  VName ->
  Param Type ->
  VName ->
  IrregularRep ->
  SubExp ->
  FlattenM (MapArray IrregularRep)
onMapIrregularInputArr :: SegLevel
-> InnerMapMode
-> Segments
-> VName
-> VName
-> VName
-> Param Type
-> VName
-> IrregularRep
-> SubExp
-> FlattenM (MapArray IrregularRep)
onMapIrregularInputArr SegLevel
lvl InnerMapMode
mode Segments
new_segments VName
ws VName
ws_O VName
ws_data Param Type
p VName
arr IrregularRep
rep SubExp
ws_prod = do
  -- new_segments already has the new w inside, unlike other functions
  rep_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType (VName -> FlattenM Type) -> VName -> FlattenM Type
forall a b. (a -> b) -> a -> b
$ IrregularRep -> VName
irregularD IrregularRep
rep
  let p_t = Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p
  when (arrayRank rep_t > 1) $
    error "onMapIrregularInputArr: irregularD is not 1D"
  if null (arrayDims p_t)
    then do
      -- Assuming irregularD is 1D, size(irregularD rep) == ws_prod should hold and this should be fine.
      let old_shape = Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
rep_t
          new_shape =
            case InnerMapMode
mode of
              InnerMapMode
SingleDim -> Segments -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
ws_prod]
              InnerMapMode
MultiDim -> Segments -> ShapeBase SubExp
segmentsShape Segments
new_segments
      case irregularK rep of
        IrregularKind
Dense -> do
          v_reshaped <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_reshaped") (Exp (Rep FlattenM) -> FlattenM VName)
-> Exp (Rep FlattenM) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$ VName -> NewShape SubExp -> BasicOp
Reshape (IrregularRep -> VName
irregularD IrregularRep
rep) (NewShape SubExp -> BasicOp) -> NewShape SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> ShapeBase SubExp -> NewShape SubExp
forall old new.
ArrayShape old =>
old -> ShapeBase new -> NewShape new
reshapeAll ShapeBase SubExp
old_shape ShapeBase SubExp
new_shape
          pure $ MapArray v_reshaped p_t
        IrregularKind
Replicated -> do
          new_flat <-
            Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName VName
arr Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_flat_expand")
              (Exp GPU -> FlattenM VName)
-> ((Solo SubExp -> FlattenM Result) -> FlattenM (Exp GPU))
-> (Solo SubExp -> FlattenM Result)
-> FlattenM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> Solo SubExp
-> (Solo SubExp -> 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 (SubExp -> Solo SubExp
forall a. a -> Solo a
MkSolo SubExp
ws_prod)
              ((Solo SubExp -> FlattenM Result) -> FlattenM VName)
-> (Solo SubExp -> FlattenM Result) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ \(MkSolo SubExp
i) -> do
                j <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"j" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ws_data [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
                data_off <- letSubExp "data_off" =<< eIndex (irregularO rep) [eSubExp j]
                seg_start <- letSubExp "seg_start" =<< eIndex ws_O [eSubExp j]
                local_pos <- letSubExp "local_pos" <=< toExp $ pe64 i - pe64 seg_start
                flat_idx <- letSubExp "flat_idx" <=< toExp $ pe64 data_off + pe64 local_pos
                fmap (subExpsRes . pure) $ letSubExp "elem" =<< eIndex (irregularD rep) [eSubExp flat_idx]
          v_reshaped <- letExp (baseName (paramName p) <> "_reshaped") $ BasicOp $ Reshape new_flat $ reshapeAll old_shape new_shape
          pure $ MapArray v_reshaped p_t
    else do
      -- We need to split multi-dimensional irregular segments into per-row
      -- segments. We compute the per-row size by dividing each segment's total
      -- size by its number of rows. The division is exact: within a single
      -- segment the array is an ordinary rectangular value, so all rows have
      -- the same size - irregularity exists only across segments. The
      -- alternative would be to read the row size from the sizes in the
      -- parameter type, but those are per-segment distributed inputs, and we do
      -- not have the environment at hand here to look them up.
      num_segments <- arraySize 0 <$> lookupType ws
      -- per_row_size[s] = irregularS[s] / ws[s]
      per_row_size <-
        letExp (baseName (paramName p) <> "_per_row_size")
          <=< segMap lvl (MkSolo num_segments)
          $ \(MkSolo SubExp
s) -> do
            total_s <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"total_s" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex (IrregularRep -> VName
irregularS IrregularRep
rep) [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
s]
            num_rows_s <- letSubExp "num_rows_s" =<< eIndex ws [eSubExp s]
            row_size <-
              letSubExp "row_size"
                =<< eIf
                  (toExp $ pe64 num_rows_s .==. 0)
                  (eBody [toExp $ intConst Int64 0])
                  (eBody [toExp $ pe64 total_s `div` pe64 num_rows_s])
            pure $ subExpsRes [row_size]
      new_S <-
        letExp (baseName (paramName p) <> "_new_S")
          <=< segMap lvl (MkSolo ws_prod)
          $ \(MkSolo SubExp
i) -> do
            seg_i <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"seg_i" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ws_data [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
            sz <- letSubExp "sz" =<< eIndex per_row_size [eSubExp seg_i]
            pure $ subExpsRes [sz]
      rep' <- case irregularK rep of
        IrregularKind
Dense -> do
          (_, new_O, m) <- SegLevel -> VName -> FlattenM (SubExp, VName, SubExp)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (SubExp, VName, SubExp)
exScanAndSum SegLevel
lvl VName
new_S
          new_F <- genFlags lvl m new_O
          pure $
            IrregularRep
              { irregularD = irregularD rep,
                irregularF = new_F,
                irregularS = new_S,
                irregularO = new_O,
                irregularK = Dense
              }
        IrregularKind
Replicated -> do
          new_O <-
            Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_new_O")
              (Exp GPU -> FlattenM VName)
-> ((Solo SubExp -> FlattenM Result) -> FlattenM (Exp GPU))
-> (Solo SubExp -> FlattenM Result)
-> FlattenM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> Solo SubExp
-> (Solo SubExp -> 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 (SubExp -> Solo SubExp
forall a. a -> Solo a
MkSolo SubExp
ws_prod)
              ((Solo SubExp -> FlattenM Result) -> FlattenM VName)
-> (Solo SubExp -> FlattenM Result) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ \(MkSolo SubExp
i) -> do
                seg_i <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"seg_i" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ws_data [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
                row_size <- letSubExp "row_size" =<< eIndex per_row_size [eSubExp seg_i]
                seg_row_start <- letSubExp "seg_row_start" =<< eIndex ws_O [eSubExp seg_i]
                row_in_seg <- letSubExp "row_in_seg" <=< toExp $ pe64 i - pe64 seg_row_start
                base_off <- letSubExp "base_off" =<< eIndex (irregularO rep) [eSubExp seg_i]
                off <- letSubExp "off" <=< toExp $ pe64 base_off + pe64 row_in_seg * pe64 row_size
                pure $ subExpsRes [off]
          m <- arraySize 0 <$> lookupType (irregularD rep)
          -- we will have mutliple write but it is the same value so it should be fine.
          new_F <- genFlags lvl m new_O
          pure $
            IrregularRep
              { irregularD = irregularD rep,
                irregularF = new_F,
                irregularS = new_S,
                irregularO = new_O,
                irregularK = Replicated
              }
      pure $ MapOther rep' p_t

onMapInputArrMultiDim ::
  SegLevel ->
  Segments ->
  SubExp ->
  DistEnv ->
  DistInputs ->
  VName ->
  VName ->
  VName ->
  Param Type ->
  VName ->
  FlattenM (MapArray IrregularRep)
onMapInputArrMultiDim :: SegLevel
-> Segments
-> SubExp
-> DistEnv
-> DistInputs
-> VName
-> VName
-> VName
-> Param Type
-> VName
-> FlattenM (MapArray IrregularRep)
onMapInputArrMultiDim SegLevel
lvl Segments
old_segments SubExp
w DistEnv
env DistInputs
inps VName
ws VName
ws_O VName
ws_data Param Type
p VName
arr = do
  case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
arr DistInputs
inps of
    Just DistInput
v_inp ->
      case DistInput
v_inp of
        DistInputFree VName
vs Type
t -> MapArray IrregularRep -> FlattenM (MapArray IrregularRep)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MapArray IrregularRep -> FlattenM (MapArray IrregularRep))
-> MapArray IrregularRep -> FlattenM (MapArray IrregularRep)
forall a b. (a -> b) -> a -> b
$ VName -> Type -> MapArray IrregularRep
forall t. VName -> Type -> MapArray t
MapArray VName
vs (Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType Type
t)
        DistInput ResTag
rt Type
t -> case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
          Irregular IrregularRep
rep -> do
            ws_prod <- Int -> Type -> SubExp
forall o. Int -> TypeBase (ShapeBase SubExp) o -> SubExp
arraySize Int
0 (Type -> SubExp) -> FlattenM Type -> FlattenM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ws_data
            onMapIrregularInputArr lvl MultiDim (old_segments <> pure w) ws ws_O ws_data p arr rep ws_prod
          Regular VName
vs -> do
            vs_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
vs
            if isAcc vs_t
              then pure $ MapArray vs (rowType t)
              else do
                -- let's be cautious and make sure it has the correct shape
                let expected_shape = Segments -> ShapeBase SubExp
segmentsShape Segments
old_segments ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t
                if arrayShape vs_t == expected_shape
                  then pure $ MapArray vs t
                  else do
                    v <-
                      letExp (baseName arr <> "_reg_reshape") . BasicOp . Reshape vs $
                        reshapeAll (arrayShape vs_t) expected_shape
                    pure $ MapArray v (rowType t)
    Maybe DistInput
Nothing -> do
      arr_row_t <- Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType (Type -> Type) -> FlattenM Type -> FlattenM Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
      arr_rep <-
        letExp (baseName arr <> "_inp_rep") . BasicOp $
          Replicate (segmentsShape old_segments) (Var arr)
      pure $ MapArray arr_rep arr_row_t

flattenMapForInBlock ::
  FlattenOps ->
  Pat Type ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM ()
flattenMapForInBlock :: FlattenOps
-> Pat Type -> SubExp -> [VName] -> Lambda SOACS -> FlattenM ()
flattenMapForInBlock FlattenOps
ops Pat Type
pat SubExp
w [VName]
arrs Lambda SOACS
map_lam = do
  scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  lam <- preprocessLambda (castScope scope) map_lam
  let arrs' = (VName -> Type -> MapArray t) -> [VName] -> [Type] -> [MapArray t]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> MapArray t
forall t. VName -> Type -> MapArray t
MapArray [VName]
arrs ([Type] -> [MapArray t]) -> [Type] -> [MapArray t]
forall a b. (a -> b) -> a -> b
$ (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam)
      (distributed, _) =
        distributeMapWith ops' scope pat [w] arrs' lam
  transformDistributed ops' mempty [w] distributed
  where
    ops' :: FlattenOps
ops' = SegLevel -> FlattenOps -> FlattenOps
atSegLevel SegLevel
inBlockSegLevel FlattenOps
ops

resultMapMode :: InnerMapMode -> DistInputs -> Type -> InnerMapMode
resultMapMode :: InnerMapMode -> DistInputs -> Type -> InnerMapMode
resultMapMode InnerMapMode
SingleDim DistInputs
_ Type
_ = InnerMapMode
SingleDim
resultMapMode InnerMapMode
MultiDim DistInputs
new_inps Type
v_t
  | (SubExp -> Bool) -> Segments -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
new_inps) (Type -> Segments
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
v_t) = InnerMapMode
SingleDim
  | Bool
otherwise = InnerMapMode
MultiDim

irregularMapResult ::
  SegLevel ->
  InnerMapMode ->
  (VName, VName, VName) ->
  Segments ->
  IrregularRep ->
  VName ->
  Type ->
  DistInputs ->
  FlattenM ResRep
irregularMapResult :: SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> Segments
-> IrregularRep
-> VName
-> Type
-> DistInputs
-> FlattenM ResRep
irregularMapResult SegLevel
lvl InnerMapMode
mode (VName
ws, VName
ws_F, VName
ws_O) Segments
segments IrregularRep
irreg VName
v Type
v_t DistInputs
new_inps =
  do
    irreg_dense <- SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_map_result") IrregularRep
irreg
    if any (isVariant new_inps) (arrayShape v_t)
      then do
        old_segment <- arraySize 0 <$> lookupType ws
        -- The size of each flattened outer segment is the sum of its rows'
        -- sizes. Because irreg_dense is dense (compact offsets), we get this in
        -- O(1) per segment as last_offset + last_size - start, avoiding a
        -- segmented reduction over the row sizes. The guard handles empty outer
        -- segments, which have no last row to read (and would index out of
        -- bounds).
        new_shape <- letExp (baseName v <> "_outer_shape") <=< segMap lvl (MkSolo old_segment) $ \(MkSolo SubExp
is) -> do
          outer_ind <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"outer_ind" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ws_O [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is]
          outer_ws_i <- letSubExp "outer_ws" =<< eIndex ws [eSubExp is]
          sz <-
            letSubExp "sz"
              =<< eIf
                (toExp $ pe64 outer_ws_i .==. 0)
                (eBody [toExp $ intConst Int64 0])
                ( do
                    last_row <- letSubExp "last_row" <=< toExp $ pe64 outer_ind + pe64 outer_ws_i - 1
                    start <- letSubExp "start" =<< eIndex (irregularO irreg_dense) [eSubExp outer_ind]
                    last_offset <- letSubExp "last_offset" =<< eIndex (irregularO irreg_dense) [eSubExp last_row]
                    last_size <- letSubExp "last_size" =<< eIndex (irregularS irreg_dense) [eSubExp last_row]
                    eBody [toExp $ pe64 last_offset - pe64 start + pe64 last_size]
                )
          pure [subExpRes sz]
        (new_ws_F, new_ws_O, _) <- doRepIota lvl new_shape
        letBindNames [v] $ BasicOp $ Replicate mempty $ Var $ irregularD irreg_dense
        mapResultRep lvl SingleDim (new_shape, new_ws_F, new_ws_O) v
      else do
        reshapeAndBind v (irregularD irreg_dense) (segmentsShape segments <> arrayShape v_t)
        mapResultRep lvl mode (ws, ws_F, ws_O) v

transformDistributedInnerMap ::
  FlattenOps ->
  InnerMapMode ->
  (VName, VName, VName) ->
  M.Map ResTag IrregularRep ->
  Segments ->
  Distributed ->
  FlattenM [(VName, ResRep)]
transformDistributedInnerMap :: FlattenOps
-> InnerMapMode
-> (VName, VName, VName)
-> Map ResTag IrregularRep
-> Segments
-> Distributed
-> FlattenM [(VName, ResRep)]
transformDistributedInnerMap FlattenOps
ops InnerMapMode
mode (VName
ws_F, VName
ws_O, VName
ws) Map ResTag IrregularRep
irregs Segments
segments Distributed
dist = do
  let Distributed DistStms
dstms (DistResults ResMap
resmap [DistRep]
reps) = Distributed
dist
  -- A name bound inside the distributed body is variant whether or not another
  -- statement uses it, so a result sized by such a name is irregular. Only
  -- counting the names that are used would leave the size existentially bound
  -- by a 'FlatMap' lambda (returned, but used by nothing) looking invariant,
  -- and the result would be given a type mentioning a name that is not in
  -- scope.
  let asInput :: DistResult -> (VName, DistInput)
asInput (DistResult ResTag
tag (DistType Segments
_ Rank
_ Type
t) VName
v) = (VName
v, ResTag -> Type -> DistInput
DistInput ResTag
tag Type
t)
      new_inps :: DistInputs
new_inps =
        (DistStm -> DistInputs) -> DistStms -> DistInputs
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DistStm -> DistInputs
distStmInputs DistStms
dstms
          DistInputs -> DistInputs -> DistInputs
forall a. Semigroup a => a -> a -> a
<> (DistResult -> (VName, DistInput)) -> [DistResult] -> DistInputs
forall a b. (a -> b) -> [a] -> [b]
map DistResult -> (VName, DistInput)
asInput ((DistStm -> [DistResult]) -> DistStms -> [DistResult]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DistStm -> [DistResult]
distStmResult DistStms
dstms)
  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_initial DistStms
dstms
  resmap_res <- fmap concat $ forM (M.toList resmap) $ \(ResTag
rt, [([DistInput], VName, Type)]
binds) ->
    [([DistInput], VName, Type)]
-> (([DistInput], VName, Type) -> FlattenM (VName, ResRep))
-> FlattenM [(VName, ResRep)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [([DistInput], VName, Type)]
binds ((([DistInput], VName, Type) -> FlattenM (VName, ResRep))
 -> FlattenM [(VName, ResRep)])
-> (([DistInput], VName, Type) -> FlattenM (VName, ResRep))
-> FlattenM [(VName, ResRep)]
forall a b. (a -> b) -> a -> b
$ \([DistInput]
cs_inps, VName
v, Type
v_t) ->
      Certs -> FlattenM (VName, ResRep) -> FlattenM (VName, ResRep)
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying (DistEnv -> [DistInput] -> Certs
distResCerts DistEnv
env [DistInput]
cs_inps) (FlattenM (VName, ResRep) -> FlattenM (VName, ResRep))
-> FlattenM (VName, ResRep) -> FlattenM (VName, ResRep)
forall a b. (a -> b) -> a -> b
$
        case (InnerMapMode -> DistInputs -> Type -> InnerMapMode
resultMapMode InnerMapMode
mode DistInputs
new_inps Type
v_t, ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env) of
          (InnerMapMode
MultiDim, Regular VName
v') ->
            if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
v_t
              then do
                [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'
                (VName, ResRep) -> FlattenM (VName, ResRep)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, VName -> ResRep
Regular VName
v)
              else do
                VName -> VName -> ShapeBase SubExp -> FlattenM ()
reshapeAndBind VName
v VName
v' (Segments -> ShapeBase SubExp
segmentsShape Segments
segments ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
v_t)
                (VName, ResRep) -> FlattenM (VName, ResRep)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, VName -> ResRep
Regular VName
v)
          (InnerMapMode
SingleDim, Regular VName
v') ->
            if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
v_t
              then do
                [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'
                (VName, ResRep) -> FlattenM (VName, ResRep)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, VName -> ResRep
Regular VName
v)
              else do
                [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'
                rep <- SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> VName
-> FlattenM ResRep
mapResultRep SegLevel
lvl InnerMapMode
SingleDim (VName
ws, VName
ws_F, VName
ws_O) VName
v
                pure (v, rep)
          (InnerMapMode
result_mode, Irregular IrregularRep
irreg) -> do
            rep <- SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> Segments
-> IrregularRep
-> VName
-> Type
-> DistInputs
-> FlattenM ResRep
irregularMapResult SegLevel
lvl InnerMapMode
result_mode (VName
ws, VName
ws_F, VName
ws_O) Segments
segments IrregularRep
irreg VName
v Type
v_t DistInputs
new_inps
            pure (v, rep)
  reps_res <- forM reps $ \(VName
v, Either SubExp DistInput
r) -> do
    case Either SubExp DistInput
r of
      Left SubExp
se -> do
        [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
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate (Segments -> ShapeBase SubExp
segmentsShape Segments
segments) SubExp
se
        -- the se is not part of input so this should be fine
        rep <- SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> VName
-> FlattenM ResRep
mapResultRep SegLevel
lvl InnerMapMode
mode (VName
ws, VName
ws_F, VName
ws_O) VName
v
        pure (v, rep)
      Right (DistInputFree VName
arr Type
t) -> do
        [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
arr
        if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
          then (VName, ResRep) -> FlattenM (VName, ResRep)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, VName -> ResRep
Regular VName
v)
          else do
            rep <- SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> VName
-> FlattenM ResRep
mapResultRep SegLevel
lvl (InnerMapMode -> DistInputs -> Type -> InnerMapMode
resultMapMode InnerMapMode
mode DistInputs
new_inps Type
t) (VName
ws, VName
ws_F, VName
ws_O) VName
v
            pure (v, rep)
      Right (DistInput ResTag
rt Type
t) ->
        let result_mode :: InnerMapMode
result_mode = InnerMapMode -> DistInputs -> Type -> InnerMapMode
resultMapMode InnerMapMode
mode DistInputs
new_inps Type
t
         in case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
              Regular VName
v' -> do
                [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'
                if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
                  then (VName, ResRep) -> FlattenM (VName, ResRep)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, VName -> ResRep
Regular VName
v)
                  else do
                    rep <- SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> VName
-> FlattenM ResRep
mapResultRep SegLevel
lvl InnerMapMode
result_mode (VName
ws, VName
ws_F, VName
ws_O) VName
v
                    pure (v, rep)
              Irregular IrregularRep
irreg -> do
                rep <- SegLevel
-> InnerMapMode
-> (VName, VName, VName)
-> Segments
-> IrregularRep
-> VName
-> Type
-> DistInputs
-> FlattenM ResRep
irregularMapResult SegLevel
lvl InnerMapMode
result_mode (VName
ws, VName
ws_F, VName
ws_O) Segments
segments IrregularRep
irreg VName
v Type
t DistInputs
new_inps
                pure (v, rep)
  pure $ resmap_res <> reps_res
  where
    env_initial :: DistEnv
env_initial = DistEnv {distResMap :: Map ResTag ResRep
distResMap = (IrregularRep -> ResRep)
-> Map ResTag IrregularRep -> Map ResTag ResRep
forall a b k. (a -> b) -> Map k a -> Map k b
M.map IrregularRep -> ResRep
Irregular Map ResTag IrregularRep
irregs}
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops

distributeAndTransformInnerMap ::
  FlattenOps ->
  InnerMapMode ->
  (VName, VName, VName) ->
  Segments ->
  DistInputs ->
  Pat Type ->
  [MapArray IrregularRep] ->
  (VName -> Maybe (FlattenM (VName, MapArray IrregularRep))) ->
  Lambda SOACS ->
  FlattenM [ResRep]
distributeAndTransformInnerMap :: FlattenOps
-> InnerMapMode
-> (VName, VName, VName)
-> Segments
-> DistInputs
-> Pat Type
-> [MapArray IrregularRep]
-> (VName -> Maybe (FlattenM (VName, MapArray IrregularRep)))
-> Lambda SOACS
-> FlattenM [ResRep]
distributeAndTransformInnerMap FlattenOps
ops InnerMapMode
mode (VName, VName, VName)
ws_triple Segments
new_segment DistInputs
inps Pat Type
pat [MapArray IrregularRep]
arrs' VName -> Maybe (FlattenM (VName, MapArray IrregularRep))
onFreeVar Lambda SOACS
map_lam = do
  -- Skip the return type: the results are described by 'pat', so variables
  -- occurring only there never need replicating into the body - and they may
  -- not even be in scope here. See Note [Ill-formed inner-map lambda].
  let free :: Names
free = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn (Lambda SOACS -> Names) -> Lambda SOACS -> Names
forall a b. (a -> b) -> a -> b
$ Lambda SOACS
map_lam {lambdaReturnType = [] :: [Type]}
  free_and_sizes <- DistInputs -> Names -> FlattenM [VName]
freeWithTypeDeps DistInputs
inps Names
free
  (free_replicated, replicated) <-
    fmap unzip . sequence $
      mapMaybe
        onFreeVar
        free_and_sizes
  free_ps <-
    zipWithM
      newParam
      (map ((<> "_free") . baseName) free_replicated) -- this should free_replicated?
      (map mapArrayRowType replicated)
  scope <- askScope
  let substs = [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$ [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
free_replicated ([VName] -> [(VName, VName)]) -> [VName] -> [(VName, VName)]
forall a b. (a -> b) -> a -> b
$ (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
free_ps
      map_lam' =
        Map VName VName -> Lambda SOACS -> Lambda SOACS
forall a. Substitute a => Map VName VName -> a -> a
substituteNames
          Map VName VName
substs
          ( Lambda SOACS
map_lam
              { lambdaParams = free_ps <> lambdaParams map_lam
              }
          )
      (distributed, arrmap) =
        distributeMapWith ops scope pat new_segment (replicated <> arrs') map_lam'
  -- order the result representations in the same order as the pattern
  resRepsInPatOrder pat
    <$> transformDistributedInnerMap ops mode ws_triple arrmap new_segment distributed

-- | Flatten a map nested in a map-nest (nonempty enclosing 'Segments'). The map
-- width is either uniform (invariant to the nest) or nonuniform, giving the
-- 'InnerMapMode': a uniform width produces a regular, multi-dimensional result
-- ('MultiDim'), while a nonuniform width is irregular and flattened into a
-- single segment dimension ('SingleDim'). The mode selects how inputs and free
-- variables are read and how the results are represented, and is threaded
-- through the rest of the flattening.
transformInnerMap ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  Pat Type ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM [ResRep]
transformInnerMap :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> Pat Type
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM [ResRep]
transformInnerMap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps Pat Type
pat SubExp
w [VName]
arrs Lambda SOACS
map_lam = do
  outer_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let mode
        | DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
w = InnerMapMode
SingleDim
        | Bool
otherwise = InnerMapMode
MultiDim
      -- In the uniform 'MultiDim' case - regular inputs and all result
      -- dimensions invariant - the flags/offsets/elements bookkeeping produced
      -- by 'doRepIota' is never consulted, so we do not emit it. This is not
      -- just an efficiency concern: when generating in-block code, the
      -- bookkeeping contains SegOps whose dimensions are bound inside the kernel
      -- body, which would make 'noNonuniformPar' reject the enclosing intrablock
      -- version.
      invariantDim Constant {} = Bool
True
      invariantDim (Var VName
v) = VName
v VName -> Scope GPU -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Scope GPU
outer_scope
      regularInput VName
arr = case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
arr DistInputs
inps of
        Just (DistInput ResTag
rt Type
_)
          | Irregular {} <- ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env -> Bool
False
        Maybe DistInput
_ -> Bool
True
      uniform =
        (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((SubExp -> Bool) -> Segments -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all SubExp -> Bool
invariantDim (Segments -> Bool) -> (Type -> Segments) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Segments
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims) (Pat Type -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat Type
pat)
          Bool -> Bool -> Bool
&& (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all VName -> Bool
regularInput [VName]
arrs
  (ws, ws_F, ws_O, ws_data) <-
    case mode of
      InnerMapMode
MultiDim
        | Bool
uniform ->
            -- XXX: this depends on laziness to explode only on usage. It might
            -- be better to handle this path more explicitly.
            (VName, VName, VName, VName)
-> FlattenM (VName, VName, VName, VName)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> VName
forall {b}. String -> b
bad String
"ws", String -> VName
forall {b}. String -> b
bad String
"ws_F", String -> VName
forall {b}. String -> b
bad String
"ws_O", String -> VName
forall {b}. String -> b
bad String
"ws_data")
      InnerMapMode
_ -> do
        ws <- SegLevel
-> Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps SubExp
w
        (ws_F, ws_O, ws_data) <- doRepIota lvl ws
        pure (ws, ws_F, ws_O, ws_data)
  (arrs', new_segment, onFreeVar) <-
    case mode of
      InnerMapMode
MultiDim -> do
        arrs' <-
          (Param Type -> VName -> FlattenM (MapArray IrregularRep))
-> [Param Type] -> [VName] -> FlattenM [MapArray IrregularRep]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM
            (SegLevel
-> Segments
-> SubExp
-> DistEnv
-> DistInputs
-> VName
-> VName
-> VName
-> Param Type
-> VName
-> FlattenM (MapArray IrregularRep)
onMapInputArrMultiDim SegLevel
lvl Segments
segments SubExp
w DistEnv
env DistInputs
inps VName
ws VName
ws_O VName
ws_data)
            (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
map_lam)
            [VName]
arrs
        pure (arrs', segments <> pure w, onMapFreeVarMultiDim lvl segments w env inps)
      InnerMapMode
SingleDim -> do
        arrs' <-
          (Param Type -> VName -> FlattenM (MapArray IrregularRep))
-> [Param Type] -> [VName] -> FlattenM [MapArray IrregularRep]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM
            (SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> VName
-> VName
-> VName
-> Param Type
-> VName
-> FlattenM (MapArray IrregularRep)
onMapInputArr SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps VName
ws VName
ws_O VName
ws_data)
            (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
map_lam)
            [VName]
arrs
        new_segment <- arraySize 0 <$> lookupType ws_data
        pure (arrs', [new_segment], onMapFreeVar lvl segments env inps ws (ws_F, ws_O, ws_data))
  distributeAndTransformInnerMap ops mode (ws_F, ws_O, ws) new_segment inps pat arrs' onFreeVar map_lam
  where
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops
    bad :: String -> b
bad String
what =
      String -> b
forall a. HasCallStack => String -> a
error (String -> b) -> String -> b
forall a b. (a -> b) -> a -> b
$ String
"transformInnerMap: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
what String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" demanded in uniform case"

-- | Flatten a map over the given enclosing 'Segments'. With no enclosing
-- segments this is a top-level map, whose inputs are ordinary regular values;
-- otherwise it is a map nested in a map-nest, whose inputs are the
-- per-enclosing-segment values.
transformMap ::
  FlattenOps ->
  -- | Incremental-flattening attributes of the map itself. These are
  -- propagated onto the (preprocessed) body, and so eventually reach maps at
  -- any depth of the nest.
  Attrs ->
  Segments ->
  DistEnv ->
  DistInputs ->
  Pat Type ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM [ResRep]
transformMap :: FlattenOps
-> Attrs
-> Segments
-> DistEnv
-> DistInputs
-> Pat Type
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM [ResRep]
transformMap FlattenOps
ops Attrs
attrs [] DistEnv
_env DistInputs
_inps Pat Type
pat SubExp
w [VName]
arrs Lambda SOACS
map_lam = do
  -- Top-level map (no enclosing segments). Preprocess the body and then
  -- propagate the attributes onto it, so they influence how the body is
  -- versioned (e.g. only_inner reaching a Screma produced by interchanging a
  -- 'sequential_outer' loop). Order matters: preprocessing may rewrite a body
  -- statement, so propagating first would lose the attributes on the rewritten
  -- form. XXX: this is arguably a bug in preprocessing.
  scope <- Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS)
-> FlattenM (Scope GPU) -> FlattenM (Scope SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope :: FlattenM (Scope SOACS)
  lam <-
    fmap (propagateVersioningAttrs attrs) . renameLambda
      =<< preprocessLambda scope map_lam
  transformTopLevelMap ops pat w arrs lam
transformMap FlattenOps
ops Attrs
attrs Segments
segments DistEnv
env DistInputs
inps Pat Type
pat SubExp
w [VName]
arrs Lambda SOACS
map_lam = do
  -- Nested map. As in the top-level case, propagate the attributes onto the
  -- preprocessed body; this is what carries them to maps deeper in the nest.
  gpu_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let pp_scope = Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS) -> Scope GPU -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ DistInputs -> Scope GPU
scopeOfDistInputs DistInputs
inps Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope GPU
gpu_scope
  lam <- propagateVersioningAttrs attrs <$> preprocessLambda pp_scope map_lam
  transformInnerMap ops segments env inps pat w arrs lam

-- | Fully flatten a map that has no enclosing segments (a top-level map). This
-- is the empty-'Segments' special case of 'transformMap': the mapped arrays are
-- ordinary regular top-level values, so we distribute the map directly over its
-- own width and flatten the resulting body, rather than reconstructing
-- per-enclosing-segment inputs. The results are necessarily regular.
transformTopLevelMap ::
  FlattenOps ->
  Pat Type ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM [ResRep]
transformTopLevelMap :: FlattenOps
-> Pat Type
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM [ResRep]
transformTopLevelMap FlattenOps
ops Pat Type
pat SubExp
w [VName]
arrs Lambda SOACS
lam = do
  scope <- Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS)
-> FlattenM (Scope GPU) -> FlattenM (Scope SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope :: FlattenM (Scope SOACS)
  let arrs' = (VName -> Type -> MapArray t) -> [VName] -> [Type] -> [MapArray t]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> MapArray t
forall t. VName -> Type -> MapArray t
MapArray [VName]
arrs ([Type] -> [MapArray t]) -> [Type] -> [MapArray t]
forall a b. (a -> b) -> a -> b
$ (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam)
      (distributed, _) =
        distributeMapWith ops scope pat [w] arrs' lam
  transformDistributed ops mempty [w] distributed
  pure $ map Regular $ patNames pat

runMapLambdaBody ::
  Segments ->
  DistEnv ->
  DistInputs ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  Pat Type ->
  [DistResult] ->
  FlattenM [VName]
runMapLambdaBody :: Segments
-> DistEnv
-> DistInputs
-> SubExp
-> [VName]
-> Lambda SOACS
-> Pat Type
-> [DistResult]
-> FlattenM [VName]
runMapLambdaBody [] DistEnv
_env DistInputs
_inps SubExp
w [VName]
arrs Lambda SOACS
map_lam Pat Type
_pat [DistResult]
_ress = do
  -- Top level (no enclosing segments): the mapped arrays are indexed directly
  -- and free variables are already in scope, so there is no per-segment input
  -- reconstruction to do - just run the (sequentialised) body under a segmap
  -- over the map width.
  map_lam' <- GLambda GPU Type -> FlattenM (GLambda GPU Type)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda (GLambda GPU Type -> FlattenM (GLambda GPU Type))
-> GLambda GPU Type -> FlattenM (GLambda GPU Type)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> GLambda GPU Type
soacsLambdaToGPU Lambda SOACS
map_lam
  vs <- letTupExp "outer_map" <=< renameExp <=< segMap defaultSegLevel [w] $ \Segments
is -> do
    let gtid :: SubExp
gtid = case Segments -> Segments
forall a. [a] -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Segments
is of
          [SubExp
i] -> SubExp
i
          Segments
_ -> String -> SubExp
forall a. HasCallStack => String -> a
error String
"runMapLambdaBody: expected single index"
    [(Param Type, VName)]
-> ((Param Type, VName) -> FlattenM ()) -> FlattenM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (GLambda GPU Type -> [LParam GPU]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda GPU Type
map_lam') [VName]
arrs) (((Param Type, VName) -> FlattenM ()) -> FlattenM ())
-> ((Param Type, VName) -> FlattenM ()) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ \(Param Type
p, VName
arr) ->
      [VName] -> Exp (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p]
        (Exp (Rep FlattenM) -> FlattenM ())
-> FlattenM (Exp (Rep FlattenM)) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< case Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p of
          Acc {} -> SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp -> FlattenM (Exp (Rep FlattenM)))
-> SubExp -> FlattenM (Exp (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
arr
          Type
_ -> VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gtid]
    GBody (Rep FlattenM) SubExpRes -> FlattenM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind (GBody (Rep FlattenM) SubExpRes -> FlattenM Result)
-> GBody (Rep FlattenM) SubExpRes -> FlattenM Result
forall a b. (a -> b) -> a -> b
$ GLambda (Rep FlattenM) Type -> GBody (Rep FlattenM) SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda (Rep FlattenM) Type
GLambda GPU Type
map_lam'
  forM vs $ \VName
v ->
    Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_copy") (Exp (Rep FlattenM) -> FlattenM VName)
-> Exp (Rep FlattenM) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate ShapeBase SubExp
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
v)
runMapLambdaBody Segments
segments DistEnv
env DistInputs
inps SubExp
w [VName]
arrs Lambda SOACS
map_lam Pat Type
_pat [DistResult]
_ress = do
  map_lam' <- GLambda GPU Type -> FlattenM (GLambda GPU Type)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda (GLambda GPU Type -> FlattenM (GLambda GPU Type))
-> GLambda GPU Type -> FlattenM (GLambda GPU Type)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> GLambda GPU Type
soacsLambdaToGPU Lambda SOACS
map_lam
  ws <- dataArr defaultSegLevel segments env inps w
  (_ws_F, ws_O, ws_data) <- doRepIota defaultSegLevel ws
  arrs' <-
    zipWithM
      (onMapInputArrMultiDim defaultSegLevel segments w env inps ws ws_O ws_data)
      (lambdaParams map_lam')
      arrs

  free_and_sizes <- freeWithTypeDeps inps (freeIn map_lam')
  (param_env, param_inputs) <-
    mapArraysToInputs (map paramName (lambdaParams map_lam')) arrs'
  let new_segments = Segments
segments Segments -> Segments -> Segments
forall a. Semigroup a => a -> a -> a
<> SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w
      free_inputs =
        [ (VName
v, DistInput
inp)
        | VName
v <- [VName]
free_and_sizes,
          Just DistInput
inp <- [VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps]
        ]

  vs <- letTupExp "outer_map" <=< renameExp <=< segMap defaultSegLevel new_segments $ \Segments
is -> do
    let full_is :: Segments
full_is = Segments -> Segments
forall a. [a] -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Segments
is
        outer_is :: Segments
outer_is = Int -> Segments -> Segments
forall a. Int -> [a] -> [a]
take (Segments -> Int
segmentsRank Segments
segments) Segments
full_is

    Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs Segments
segments DistEnv
env Segments
outer_is DistInputs
free_inputs
    Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs Segments
new_segments DistEnv
param_env Segments
full_is DistInputs
param_inputs

    GBody (Rep FlattenM) SubExpRes -> FlattenM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind (GBody (Rep FlattenM) SubExpRes -> FlattenM Result)
-> GBody (Rep FlattenM) SubExpRes -> FlattenM Result
forall a b. (a -> b) -> a -> b
$ GLambda (Rep FlattenM) Type -> GBody (Rep FlattenM) SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda (Rep FlattenM) Type
GLambda GPU Type
map_lam'
  forM vs $ \VName
v -> do
    Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_copy") (Exp (Rep FlattenM) -> FlattenM VName)
-> Exp (Rep FlattenM) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$
        ShapeBase SubExp -> SubExp -> BasicOp
Replicate ShapeBase SubExp
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
v)

versionedUniformMap ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Pat Type ->
  StmAux () ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM DistEnv
versionedUniformMap :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM DistEnv
versionedUniformMap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
ress Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs Lambda SOACS
map_lam = do
  let only_intra :: Bool
only_intra = Attrs -> Bool
onlyExploitIntra (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
aux)
      may_intra :: Bool
may_intra = Lambda SOACS -> Bool
worthIntrablock Lambda SOACS
map_lam Bool -> Bool -> Bool
&& Attrs -> Bool
mayExploitIntra (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
aux)

  intra' <-
    if Bool
only_intra Bool -> Bool -> Bool
|| Bool
may_intra
      then (Pat Type -> SubExp -> [VName] -> Lambda SOACS -> FlattenM ())
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM (Maybe IntrablockResult)
Intrablock.intrablockParallelise (FlattenOps
-> Pat Type -> SubExp -> [VName] -> Lambda SOACS -> FlattenM ()
flattenMapForInBlock FlattenOps
ops) Segments
segments DistEnv
env DistInputs
inps [DistResult]
ress Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs Lambda SOACS
map_lam
      else Maybe IntrablockResult -> FlattenM (Maybe IntrablockResult)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe IntrablockResult
forall a. Maybe a
Nothing

  let fullFlatten =
        [ResRep] -> [VName]
regularRepVars ([ResRep] -> [VName]) -> FlattenM [ResRep] -> FlattenM [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FlattenOps
-> Attrs
-> Segments
-> DistEnv
-> DistInputs
-> Pat Type
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM [ResRep]
transformMap (SegLevel -> FlattenOps -> FlattenOps
atSegLevel SegLevel
defaultSegLevel FlattenOps
ops) (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
aux) Segments
segments DistEnv
env DistInputs
inps Pat Type
pat SubExp
w [VName]
arrs Lambda SOACS
map_lam

      outerOnly =
        Segments
-> DistEnv
-> DistInputs
-> SubExp
-> [VName]
-> Lambda SOACS
-> Pat Type
-> [DistResult]
-> FlattenM [VName]
runMapLambdaBody Segments
segments DistEnv
env DistInputs
inps SubExp
w [VName]
arrs Lambda SOACS
map_lam Pat Type
pat [DistResult]
ress

      result_ts =
        [ Type
t Type -> ShapeBase SubExp -> Type
`arrayOfShape` Segments -> ShapeBase SubExp
segmentsShape Segments
segments
        | DistResult ResTag
_ (DistType Segments
_ Rank
_ Type
t) VName
_ <- [DistResult]
ress
        ]

  match_res <-
    certifying (distCerts inps aux env) $
      mapAlternatives
        "match_res"
        result_ts
        (stmAuxAttrs aux)
        -- 'versionedUniformMap' is only reached via 'isVersionableMap', which
        -- guarantees the body calls no parallel function.
        False
        (worthSequentialising map_lam)
        (segments <> pure w)
        (renamedBody fullFlatten)
        (renamedBody outerOnly)
        intra'

  pure $ insertRegulars (map distResTag ress) match_res env

flattenUniformRedomap ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Pat Type ->
  StmAux () ->
  SubExp ->
  [VName] ->
  ScremaForm SOACS ->
  [Reduce SOACS] ->
  Lambda SOACS ->
  FlattenM DistEnv
flattenUniformRedomap :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> [Reduce SOACS]
-> Lambda SOACS
-> FlattenM DistEnv
flattenUniformRedomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs ScremaForm SOACS
form [Reduce SOACS]
reds Lambda SOACS
map_lam = do
  let outer_only :: FlattenM [VName]
outer_only = SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> SubExp
-> [VName]
-> [Reduce SOACS]
-> Lambda SOACS
-> FlattenM [VName]
transformUniformRedomap (FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops) Segments
segments DistEnv
env DistInputs
inps SubExp
w [VName]
arrs [Reduce SOACS]
reds Lambda SOACS
map_lam
  gpu_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let pp_scope = Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS) -> Scope GPU -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ DistInputs -> Scope GPU
scopeOfDistInputs DistInputs
inps Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope GPU
gpu_scope
  factored <- factorScremaForParallelism funHasParallelism pp_scope aux pat w arrs form
  case factored of
    Just Body SOACS
body ->
      FlattenOps
-> Name
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> SubExp
-> Body SOACS
-> FlattenM [VName]
-> FlattenM DistEnv
versionScanRed FlattenOps
ops Name
"uniform_redomap_alt" Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux SubExp
w Body SOACS
body FlattenM [VName]
outer_only
    Maybe (Body SOACS)
Nothing -> do
      elems' <- FlattenM [VName]
outer_only
      pure $ insertRegulars (map distResTag res) elems' env
  where
    funHasParallelism :: FunHasParallelism
funHasParallelism = FlattenOps -> FunHasParallelism
flattenFunHasParallelism FlattenOps
ops

flattenSegRedomap ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  SubExp ->
  [VName] ->
  [Reduce SOACS] ->
  Lambda SOACS ->
  FlattenM DistEnv
flattenSegRedomap :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> SubExp
-> [VName]
-> [Reduce SOACS]
-> Lambda SOACS
-> FlattenM DistEnv
flattenSegRedomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res SubExp
w [VName]
arrs [Reduce SOACS]
reds Lambda SOACS
map_lam = do
  reps <- (VName -> FlattenM ResRep) -> [VName] -> FlattenM [ResRep]
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 (SegLevel
-> Segments -> DistEnv -> DistInputs -> VName -> FlattenM ResRep
segOpInputRep SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps) [VName]
arrs
  let sing_red = [Reduce SOACS] -> Reduce SOACS
forall rep. Buildable rep => [Reduce rep] -> Reduce rep
singleReduce [Reduce SOACS]
reds
      hasNoFreeVariant = (VName -> Bool) -> Names -> Bool
allNames (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (SubExp -> Bool) -> (VName -> SubExp) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) (Reduce SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Reduce SOACS
sing_red Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
map_lam)
  (ws_F, ws_O, ws_S, elems, elems_kind) <-
    prepareSegOpInputs lvl segments env inps w reps arrs hasNoFreeVariant
  nes' <- mapM (readNeutral segments env inps) (redNeutral sing_red)
  let sing_red' = Reduce SOACS
sing_red {redNeutral = nes'}
  let free = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
map_lam
  free_and_sizes <- freeWithTypeDeps inps free
  ws <- dataArr lvl segments env inps w
  (_, _, ws_data) <- doRepIota lvl ws_S
  (free_replicated, replicated) <-
    fmap unzip . sequence $
      mapMaybe
        (onMapFreeVar lvl segments env inps ws (ws_F, ws_O, ws_data))
        free_and_sizes
  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated

  new_segment <- arraySize 0 <$> lookupType ws_F
  let readFree Segments
is = Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs [SubExp
new_segment] DistEnv
free_env Segments
is DistInputs
free_inputs
  (red_elems, mapout_elems) <-
    genSegRedomap lvl ws_S ws_F ws_O elems sing_red' (soacsLambdaToGPU map_lam) readFree
  red_elems' <- forM red_elems $ \VName
v -> do
    v_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
    letExp (baseName v <> "_reshaped") . BasicOp $
      Reshape v $
        reshapeAll (arrayShape v_t) (segmentsShape segments)
  let (red_res, map_res) = splitAt (redResults reds) res
  env' <-
    insertSegOpMapResults
      segments
      ws_S
      ws_F
      ws_O
      elems_kind
      (zip map_res mapout_elems)
      env
  pure $ insertRegulars (map distResTag red_res) red_elems' env'
  where
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops

flattenUniformMaposcanomap ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Pat Type ->
  StmAux () ->
  SubExp ->
  [VName] ->
  ScremaForm SOACS ->
  [Scan SOACS] ->
  Lambda SOACS ->
  Lambda SOACS ->
  FlattenM DistEnv
flattenUniformMaposcanomap :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> [Scan SOACS]
-> Lambda SOACS
-> Lambda SOACS
-> FlattenM DistEnv
flattenUniformMaposcanomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs ScremaForm SOACS
form [Scan SOACS]
scans Lambda SOACS
post_lam Lambda SOACS
map_lam = do
  let outer_only :: FlattenM [VName]
outer_only =
        SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> SubExp
-> [VName]
-> [Scan SOACS]
-> Lambda SOACS
-> Lambda SOACS
-> FlattenM [VName]
transformUniformMaposcanomap SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps SubExp
w [VName]
arrs [Scan SOACS]
scans Lambda SOACS
post_lam Lambda SOACS
map_lam
  gpu_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let pp_scope = Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS) -> Scope GPU -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ DistInputs -> Scope GPU
scopeOfDistInputs DistInputs
inps Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope GPU
gpu_scope
  factored <- factorScremaForParallelism funHasParallelism pp_scope aux pat w arrs form
  case factored of
    Just Body SOACS
body ->
      FlattenOps
-> Name
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> SubExp
-> Body SOACS
-> FlattenM [VName]
-> FlattenM DistEnv
versionScanRed FlattenOps
ops Name
"uniform_maposcanomap_alt" Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux SubExp
w Body SOACS
body FlattenM [VName]
outer_only
    Maybe (Body SOACS)
Nothing -> do
      elems' <- FlattenM [VName]
outer_only
      pure $ insertRegulars (map distResTag res) elems' env
  where
    funHasParallelism :: FunHasParallelism
funHasParallelism = FlattenOps -> FunHasParallelism
flattenFunHasParallelism FlattenOps
ops
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops

flattenSegMaposcanomap ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  SubExp ->
  [VName] ->
  [Scan SOACS] ->
  Lambda SOACS ->
  Lambda SOACS ->
  FlattenM DistEnv
flattenSegMaposcanomap :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> SubExp
-> [VName]
-> [Scan SOACS]
-> Lambda SOACS
-> Lambda SOACS
-> FlattenM DistEnv
flattenSegMaposcanomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res SubExp
w [VName]
arrs [Scan SOACS]
scans Lambda SOACS
post_lam Lambda SOACS
map_lam = do
  reps <- (VName -> FlattenM ResRep) -> [VName] -> FlattenM [ResRep]
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 (SegLevel
-> Segments -> DistEnv -> DistInputs -> VName -> FlattenM ResRep
segOpInputRep SegLevel
lvl Segments
segments DistEnv
env DistInputs
inps) [VName]
arrs
  let hasNoFreeVariant = (VName -> Bool) -> Names -> Bool
allNames (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (SubExp -> Bool) -> (VName -> SubExp) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) (Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
post_lam Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
map_lam Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> (Scan SOACS -> Names) -> [Scan SOACS] -> Names
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Scan SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn [Scan SOACS]
scans)
  (ws_F, ws_O, ws_S, elems, elems_kind) <-
    prepareSegOpInputs lvl segments env inps w reps arrs hasNoFreeVariant
  let free = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
map_lam Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
post_lam
  free_and_sizes <- freeWithTypeDeps inps free
  ws <- dataArr lvl segments env inps w
  (_, _, ws_data) <- doRepIota lvl ws_S
  (free_replicated, replicated) <-
    fmap unzip . sequence $
      mapMaybe
        (onMapFreeVar lvl segments env inps ws (ws_F, ws_O, ws_data))
        free_and_sizes
  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated
  new_segment <- arraySize 0 <$> lookupType ws_F
  let readFree Segments
is = Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs [SubExp
new_segment] DistEnv
free_env Segments
is DistInputs
free_inputs
  elems' <- doSegMaposcanomap lvl scans ws_F elems post_lam map_lam segments inps env readFree
  insertSegOpMapResults
    segments
    ws_S
    ws_F
    ws_O
    elems_kind
    (zip res elems')
    env
  where
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops

flattenPlainMap ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Pat Type ->
  StmAux () ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM DistEnv
flattenPlainMap :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM DistEnv
flattenPlainMap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs Lambda SOACS
map_lam = do
  map_res <-
    FlattenOps
-> Attrs
-> Segments
-> DistEnv
-> DistInputs
-> Pat Type
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM [ResRep]
transformMap FlattenOps
ops (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
aux) Segments
segments DistEnv
env DistInputs
inps Pat Type
pat SubExp
w [VName]
arrs Lambda SOACS
map_lam
  insertRepsM (zip (map distResTag res) map_res) env

flattenOtherScrema ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Pat Type ->
  StmAux () ->
  SubExp ->
  [VName] ->
  ScremaForm SOACS ->
  FlattenM DistEnv
flattenOtherScrema :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> FlattenM DistEnv
flattenOtherScrema FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs ScremaForm SOACS
form = do
  gpu_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let pp_scope = Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS) -> Scope GPU -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ DistInputs -> Scope GPU
scopeOfDistInputs DistInputs
inps Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope GPU
gpu_scope
  factored <- factorScremaForParallelism funHasParallelism pp_scope aux pat w arrs form
  case factored of
    Just Body SOACS
body -> do
      reps <- FlattenOps
-> Segments
-> Name
-> DistEnv
-> DistInputs
-> [DistResult]
-> Body SOACS
-> FlattenM [ResRep]
distributeAndFlattenBody FlattenOps
ops Segments
segments Name
"factorScremaForParallelism_body" DistEnv
env DistInputs
inps [DistResult]
res Body SOACS
body
      insertRepsM (zip (map distResTag res) reps) env
    Maybe (Body SOACS)
Nothing
      -- XXX: here we silently sequentialise any SOAC that is not handled
      -- above if it is possible to do so. We need to make sure that we
      -- actually handle everything we care about!
      | ScremaForm SOACS -> Bool
shouldDissectForm ScremaForm SOACS
form ->
          String -> FlattenM DistEnv
forall a. HasCallStack => String -> a
error String
"flattenScrema: complex Screma survived preprocessing"
      | (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
res ->
          FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
flattenScalarStm FlattenOps
ops 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 (Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))
      | Bool
otherwise -> do
          -- XXX: The results are nonuniform, so we cannot run the SOAC
          -- unchanged inside a kernel. Sequentialise it to a loop and
          -- flatten that instead. This does lose us potential parallelism.
          -- A solution would be to preprocess such cases to express them in
          -- terms of loops and maps instead, which we can indeed handle.
          stms <-
            Scope SOACS -> Seq (Stm SOACS) -> FlattenM (Seq (Stm SOACS))
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Seq (Stm SOACS) -> m (Seq (Stm SOACS))
preprocessStms Scope SOACS
pp_scope
              (Seq (Stm SOACS) -> FlattenM (Seq (Stm SOACS)))
-> FlattenM (Seq (Stm SOACS)) -> FlattenM (Seq (Stm SOACS))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Scope SOACS
-> BuilderT SOACS FlattenM () -> FlattenM (Seq (Stm SOACS))
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> BuilderT SOACS m a -> m (Seq (Stm SOACS))
runSimplifiedBuilder
                Scope SOACS
pp_scope
                (StmAux ()
-> BuilderT SOACS FlattenM () -> BuilderT SOACS FlattenM ()
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
aux (BuilderT SOACS FlattenM () -> BuilderT SOACS FlattenM ())
-> BuilderT SOACS FlattenM () -> BuilderT SOACS FlattenM ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (BuilderT SOACS FlattenM)))
-> SOAC (Rep (BuilderT SOACS FlattenM))
-> BuilderT SOACS FlattenM ()
forall (m :: * -> *).
Transformer m =>
Pat (LetDec (Rep m)) -> SOAC (Rep m) -> m ()
FOT.transformSOAC Pat Type
Pat (LetDec (Rep (BuilderT SOACS FlattenM)))
pat (SOAC (Rep (BuilderT SOACS FlattenM))
 -> BuilderT SOACS FlattenM ())
-> SOAC (Rep (BuilderT SOACS FlattenM))
-> BuilderT SOACS FlattenM ()
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName]
-> ScremaForm (Rep (BuilderT SOACS FlattenM))
-> SOAC (Rep (BuilderT SOACS FlattenM))
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs ScremaForm (Rep (BuilderT SOACS FlattenM))
ScremaForm SOACS
form)
          let body = Seq (Stm SOACS) -> Result -> Body SOACS
forall res.
IsResult res =>
Seq (Stm SOACS) -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Seq (Stm SOACS)
stms (Result -> Body SOACS) -> Result -> Body SOACS
forall a b. (a -> b) -> a -> b
$ [VName] -> Result
varsRes ([VName] -> Result) -> [VName] -> Result
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pat
          reps <- distributeAndFlattenBody ops segments "sequentialised_soac" env inps res body
          insertRepsM (zip (map distResTag res) reps) env
  where
    funHasParallelism :: FunHasParallelism
funHasParallelism = FlattenOps -> FunHasParallelism
flattenFunHasParallelism FlattenOps
ops

flattenScrema ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  (Pat Type, StmAux ()) ->
  (SubExp, [VName], ScremaForm SOACS) ->
  FlattenM DistEnv
flattenScrema :: 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, StmAux ()
aux) (SubExp
w, [VName]
arrs, ScremaForm SOACS
form)
  | Just ([Reduce SOACS]
reds, Lambda SOACS
map_lam) <- ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm SOACS
form,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
w,
    (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
res,
    (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (DistEnv -> DistInputs -> VName -> Bool
isRegularInputArr DistEnv
env DistInputs
inps) [VName]
arrs,
    (Reduce SOACS -> Bool) -> [Reduce SOACS] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (DistInputs -> Lambda SOACS -> Bool
suitableUniformOperator DistInputs
inps (Lambda SOACS -> Bool)
-> (Reduce SOACS -> Lambda SOACS) -> Reduce SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda) [Reduce SOACS]
reds =
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> [Reduce SOACS]
-> Lambda SOACS
-> FlattenM DistEnv
flattenUniformRedomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs ScremaForm SOACS
form [Reduce SOACS]
reds Lambda SOACS
map_lam
  | Just ([Reduce SOACS]
reds, Lambda SOACS
map_lam) <- ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm SOACS
form,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasParallelism FunHasParallelism
funHasParallelism Lambda SOACS
map_lam,
    (Reduce SOACS -> Bool) -> [Reduce SOACS] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Reduce SOACS
red -> DistEnv -> DistInputs -> Lambda SOACS -> Segments -> Bool
suitableOperator DistEnv
env DistInputs
inps (Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda Reduce SOACS
red) (Reduce SOACS -> Segments
forall rep. Reduce rep -> Segments
redNeutral Reduce SOACS
red)) [Reduce SOACS]
reds =
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> SubExp
-> [VName]
-> [Reduce SOACS]
-> Lambda SOACS
-> FlattenM DistEnv
flattenSegRedomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res SubExp
w [VName]
arrs [Reduce SOACS]
reds Lambda SOACS
map_lam
  | Just (Lambda SOACS
post_lam, [Scan SOACS]
scans, Lambda SOACS
map_lam) <- ScremaForm SOACS
-> Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS)
forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC ScremaForm SOACS
form,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
w,
    (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
res,
    (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (DistEnv -> DistInputs -> VName -> Bool
isRegularInputArr DistEnv
env DistInputs
inps) [VName]
arrs,
    (Scan SOACS -> Bool) -> [Scan SOACS] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (DistInputs -> Lambda SOACS -> Bool
suitableUniformOperator DistInputs
inps (Lambda SOACS -> Bool)
-> (Scan SOACS -> Lambda SOACS) -> Scan SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scan SOACS -> Lambda SOACS
forall rep. Scan rep -> Lambda rep
scanLambda) [Scan SOACS]
scans =
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> [Scan SOACS]
-> Lambda SOACS
-> Lambda SOACS
-> FlattenM DistEnv
flattenUniformMaposcanomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs ScremaForm SOACS
form [Scan SOACS]
scans Lambda SOACS
post_lam Lambda SOACS
map_lam
  | Just (Lambda SOACS
post_lam, [Scan SOACS]
scans, Lambda SOACS
map_lam) <- ScremaForm SOACS
-> Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS)
forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC ScremaForm SOACS
form,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasParallelism FunHasParallelism
funHasParallelism Lambda SOACS
map_lam,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ FunHasParallelism -> Lambda SOACS -> Bool
lambdaHasParallelism FunHasParallelism
funHasParallelism Lambda SOACS
post_lam,
    (Scan SOACS -> Bool) -> [Scan SOACS] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Scan SOACS
scan -> DistEnv -> DistInputs -> Lambda SOACS -> Segments -> Bool
suitableOperator DistEnv
env DistInputs
inps (Scan SOACS -> Lambda SOACS
forall rep. Scan rep -> Lambda rep
scanLambda Scan SOACS
scan) (Scan SOACS -> Segments
forall rep. Scan rep -> Segments
scanNeutral Scan SOACS
scan)) [Scan SOACS]
scans =
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> SubExp
-> [VName]
-> [Scan SOACS]
-> Lambda SOACS
-> Lambda SOACS
-> FlattenM DistEnv
flattenSegMaposcanomap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res SubExp
w [VName]
arrs [Scan SOACS]
scans Lambda SOACS
post_lam Lambda SOACS
map_lam
  | Just Lambda SOACS
map_lam <- ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form,
    FunHasParallelism
-> SegLevel
-> DistInputs
-> DistEnv
-> SubExp
-> [DistResult]
-> Lambda SOACS
-> Bool
isVersionableMap FunHasParallelism
funHasParallelism SegLevel
lvl DistInputs
inps DistEnv
env SubExp
w [DistResult]
res Lambda SOACS
map_lam =
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM DistEnv
versionedUniformMap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs Lambda SOACS
map_lam
  | Just Lambda SOACS
map_lam <- ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form =
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM DistEnv
flattenPlainMap FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs Lambda SOACS
map_lam
  | Bool
otherwise =
      FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> ScremaForm SOACS
-> FlattenM DistEnv
flattenOtherScrema FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res Pat Type
pat StmAux ()
aux SubExp
w [VName]
arrs ScremaForm SOACS
form
  where
    funHasParallelism :: FunHasParallelism
funHasParallelism = FlattenOps -> FunHasParallelism
flattenFunHasParallelism FlattenOps
ops
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops

-- Note [FlatMap element counting]
--
-- 'FlatMap' counts in units of the lambda's element type @b@: the data array
-- for irregular results is @[m]b@, the shape and offset arrays hold per-segment
-- counts of @b@-elements, and the flag array has one entry per @b@-element.
--
-- The flattening pass, however, represents irregular data with an
-- 'IrregularRep' whose 'irregularD' is primitive (a flat array of scalars), so
-- its segment sizes and flags count *scalars*. When @b@ is itself a @c@-element
-- array there are @c@ scalars per @b@-element ('flatMapElemsPer'), so the two
-- notions of "size" differ by a factor of @c@.
--
-- The flattening rules therefore convert between the two: the metadata that
-- 'FlatMap' surfaces directly (the total @m@, and the shape, flag, and offset
-- arrays) is put in @b@-units by dividing scalar-unit sizes by @c@. The data
-- arrays themselves, however, keep their natural primitive 'IrregularRep' with
-- scalar-unit structure - the "'irregularD' is always primitive" invariant is
-- preserved. A consumer of such a data array reconstructs the @b@-element rows
-- from the primitive rep: 'onMapIrregularInputArr' does so for a plain map, and
-- 'prepareSegOpInputs' for a segmented redomap/scan.

-- Note [Ill-formed inner-map lambda]
--
-- 'flattenFlatMapNested' hands the inner-map machinery ('transformMap') an
-- ordinary 'Lambda' obtained from the 'FlatMap' 'ExtLambda' by instantiating
-- the existential size with the lambda's own leading result. That lambda is
-- ill-formed: the size is bound by a statement in the body, and so is
-- meaningless in the return type, which describes the lambda from outside.
--
-- We do it because 'transformMap' requires a 'Lambda', whose types cannot
-- express a result whose outer size varies per iteration, which is exactly what
-- an irregular 'FlatMap' result is.
--
-- Nothing downstream derives meaning from the return type - the results are
-- described by 'map_pat' throughout - but it is still traversed in two places,
-- which is what makes this contained rather than harmless:
--
-- - 'distributeAndTransformInnerMap' computes the free variables that must be
--   replicated into the body, and must skip the return type; otherwise the size
--   is looked up in the enclosing scope, where it does not exist.
--
-- - 'preprocessLambda' simplifies the return type outside the scope of the
--   parameters, where a size that resolves to nothing is left alone.
--
-- FIXME: Teaching the inner-map machinery to accept an 'ExtLambda' would remove
-- the need for all of this.

-- | The number of scalars per element of a 'FlatMap' lambda's (first) value
-- result; see Note [FlatMap element counting].
flatMapElemsPer :: ExtLambda SOACS -> TPrimExp Int64 VName
flatMapElemsPer :: ExtLambda SOACS -> TPrimExp Int64 VName
flatMapElemsPer ExtLambda SOACS
lam =
  [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ (SubExp -> TPrimExp Int64 VName)
-> Segments -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 (Segments -> [TPrimExp Int64 VName])
-> Segments -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ Type -> Segments
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> Segments) -> Type -> Segments
forall a b. (a -> b) -> a -> b
$ [Type] -> Type
forall a. HasCallStack => [a] -> a
head ([Type] -> Type) -> [Type] -> Type
forall a b. (a -> b) -> a -> b
$ ExtLambda SOACS -> [Type]
forall rep. ExtLambda rep -> [Type]
flatMapRowTypes ExtLambda SOACS
lam

-- | Divide each entry of a per-segment size (or offset) array by the given
-- factor, converting scalar-unit counts to element-unit counts. A no-op when the
-- factor is statically 1 (a scalar element type). See Note [FlatMap element
-- counting].
scaleSizesDown :: SegLevel -> TPrimExp Int64 VName -> VName -> FlattenM VName
scaleSizesDown :: SegLevel -> TPrimExp Int64 VName -> VName -> FlattenM VName
scaleSizesDown SegLevel
_ TPrimExp Int64 VName
elems_per VName
arr | TPrimExp Int64 VName
elems_per TPrimExp Int64 VName -> TPrimExp Int64 VName -> Bool
forall a. Eq a => a -> a -> Bool
== TPrimExp Int64 VName
1 = VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
arr
scaleSizesDown SegLevel
lvl TPrimExp Int64 VName
elems_per VName
arr = do
  arr_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
  letExp "scaled_sizes" <=< segMap lvl (arrayDims arr_t) $ \Segments
gtids -> do
    x <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"x" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr ((SubExp -> FlattenM (Exp GPU)) -> Segments -> [FlattenM (Exp GPU)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> FlattenM (Exp (Rep FlattenM))
SubExp -> FlattenM (Exp GPU)
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp Segments
gtids)
    x_b <- letSubExp "x_b" =<< toExp (pe64 x `div` elems_per)
    pure [subExpRes x_b]

-- | Subsample a scalar-unit flag array to element units by taking every
-- @elems_per@-th entry, giving one flag per element. A no-op when the factor is
-- statically 1. This reuses the (already-built) flag array rather than
-- recomputing segment starts with a segmented scan. See Note [FlatMap element
-- counting].
subsampleFlags :: SegLevel -> TPrimExp Int64 VName -> VName -> FlattenM VName
subsampleFlags :: SegLevel -> TPrimExp Int64 VName -> VName -> FlattenM VName
subsampleFlags SegLevel
_ TPrimExp Int64 VName
elems_per VName
arr | TPrimExp Int64 VName
elems_per TPrimExp Int64 VName -> TPrimExp Int64 VName -> Bool
forall a. Eq a => a -> a -> Bool
== TPrimExp Int64 VName
1 = VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
arr
subsampleFlags SegLevel
lvl TPrimExp Int64 VName
elems_per VName
arr = do
  big_m <- Int -> Type -> SubExp
forall o. Int -> TypeBase (ShapeBase SubExp) o -> SubExp
arraySize Int
0 (Type -> SubExp) -> FlattenM Type -> FlattenM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
  m <- letSubExp "flags_m" =<< toExp (pe64 big_m `div` elems_per)
  letExp "elem_flags" <=< segMap lvl (MkSolo m) $ \(MkSolo SubExp
i) -> do
    flag <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"flag" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr [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 (TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM)))
-> TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
i TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per]
    pure [subExpRes flag]

-- | Flattening rule for a top-level 'FlatMap', which is a very thin wrapper
-- over just flattening the lambda. The result produced by 'FlatMap' corresponds
-- exactly to the internal irregular representation, so we simply distribute the
-- lambda over the @w@ segments and obtain an 'IrregularRep' for each result
-- (they share the same segment structure, as required by the type of
-- 'FlatMap').
flattenFlatMap ::
  FlattenOps ->
  Pat Type ->
  SubExp ->
  [VName] ->
  ExtLambda SOACS ->
  FlattenM ()
flattenFlatMap :: FlattenOps
-> Pat Type -> SubExp -> [VName] -> ExtLambda SOACS -> FlattenM ()
flattenFlatMap FlattenOps
ops Pat Type
pat SubExp
w [VName]
arrs ExtLambda SOACS
lam = do
  let segments :: Segments
segments = [SubExp
w]
      inps :: DistInputs
inps =
        (Param Type -> VName -> (VName, DistInput))
-> [Param Type] -> [VName] -> DistInputs
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          (\Param Type
p VName
arr -> (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p, VName -> Type -> DistInput
DistInputFree VName
arr (Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p)))
          (ExtLambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams ExtLambda SOACS
lam)
          [VName]
arrs
      elem_ts :: [Type]
elem_ts = ExtLambda SOACS -> [Type]
forall rep. ExtLambda rep -> [Type]
flatMapRowTypes ExtLambda SOACS
lam
      (VName
m_name, VName
s_name, VName
f_name, VName
o_name, [VName]
value_names) =
        case Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pat of
          (VName
a : VName
b : VName
c : VName
d : [VName]
hs) -> (VName
a, VName
b, VName
c, VName
d, [VName]
hs)
          [VName]
_ -> String -> (VName, VName, VName, VName, [VName])
forall a. HasCallStack => String -> a
error String
"flattenFlatMap: pattern too short"
      ([VName]
d_names, [VName]
r_names) = ExtLambda SOACS -> [VName] -> ([VName], [VName])
forall rep a. ExtLambda rep -> [a] -> ([a], [a])
flatMapSplitValues ExtLambda SOACS
lam [VName]
value_names
      -- A nonuniform result is distributed as a variably sized array, a uniform
      -- one as an ordinary array with one element per segment.
      distTypeOf :: ExtType -> DistType
distTypeOf ExtType
t
        | ExtType -> Bool
flatMapNonuniform ExtType
t = Segments -> Rank -> Type -> DistType
DistType Segments
segments (Int -> Rank
Rank Int
1) (Type -> DistType) -> Type -> DistType
forall a b. (a -> b) -> a -> b
$ ExtType -> Type
forall {o}.
TypeBase (ShapeBase ExtSize) o -> TypeBase (ShapeBase SubExp) o
static (ExtType -> Type) -> ExtType -> Type
forall a b. (a -> b) -> a -> b
$ ExtType -> ExtType
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType ExtType
t
        | Bool
otherwise = Segments -> Rank -> Type -> DistType
DistType Segments
segments (Int -> Rank
Rank Int
0) (Type -> DistType) -> Type -> DistType
forall a b. (a -> b) -> a -> b
$ ExtType -> Type
forall {o}.
TypeBase (ShapeBase ExtSize) o -> TypeBase (ShapeBase SubExp) o
static ExtType
t
      res :: [DistResult]
res =
        (Int -> VName -> ExtType -> DistResult)
-> [Int] -> [VName] -> [ExtType] -> [DistResult]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3
          (\Int
i VName
v ExtType
t -> ResTag -> DistType -> VName -> DistResult
DistResult (Int -> ResTag
ResTag Int
i) (ExtType -> DistType
distTypeOf ExtType
t) VName
v)
          [Int
0 ..]
          [VName]
value_names
          [ExtType]
value_ts
      -- The lambda's leading size result is of no use here, as the metadata is
      -- derived from the segment structure of the flattened irregular results.
      body :: Body SOACS
body = ExtLambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody ExtLambda SOACS
lam
      body' :: Body SOACS
body' = Body SOACS
body {bodyResult = drop 1 (bodyResult body)}
  reps <- FlattenOps
-> Segments
-> Name
-> DistEnv
-> DistInputs
-> [DistResult]
-> Body SOACS
-> FlattenM [ResRep]
distributeAndFlattenBody FlattenOps
ops Segments
segments Name
"flatmap" DistEnv
forall a. Monoid a => a
mempty DistInputs
inps [DistResult]
res Body SOACS
body'
  let (irreg_reps, reg_reps) = flatMapSplitValues lam reps
  irregs <- forM irreg_reps $ \case
    Irregular IrregularRep
ir -> SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl Name
"flatmap_res" IrregularRep
ir
    Regular VName
_ -> String -> FlattenM IrregularRep
forall a. HasCallStack => String -> a
error String
"flattenFlatMap: irregular result is not irregular"
  -- The regular results (outer size @w@) are already in the form we want.
  forM_ (zip r_names reg_reps) $ \(VName
v, ResRep
rep) -> case ResRep
rep of
    Regular 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'
    Irregular IrregularRep
_ -> String -> FlattenM ()
forall a. HasCallStack => String -> a
error String
"flattenFlatMap: regular result is not regular"
  case (irregs, elem_ts) of
    (IrregularRep
ir0 : [IrregularRep]
_, Type
_ : [Type]
_) -> do
      -- The flattening structure arrays are in units of scalars, but the source
      -- 'flatmap' wants them in units of the lambda's element type @b@, of
      -- which there are 'elems_per' scalars. All results share the same segment
      -- structure, so we derive the source metadata once, from the first
      -- result:
      --
      --   * The shape and offset arrays (per-segment, size @w@) are the
      --     scalar-unit 'irregularS' and 'irregularO' scaled down by
      --     'elems_per'.
      --
      --   * the flag array (per-element, size @m@) is 'irregularF' subsampled
      --     with stride 'elems_per'.
      let elems_per :: TPrimExp Int64 VName
elems_per = ExtLambda SOACS -> TPrimExp Int64 VName
flatMapElemsPer ExtLambda SOACS
lam
      -- The total size is the number of data elements.
      big_m <- Int -> Type -> SubExp
forall o. Int -> TypeBase (ShapeBase SubExp) o -> SubExp
arraySize Int
0 (Type -> SubExp) -> FlattenM Type -> FlattenM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType (IrregularRep -> VName
irregularD IrregularRep
ir0)
      m <- letSubExp "flatmap_m" =<< toExp (pe64 big_m `div` elems_per)
      letBindNames [m_name] $ BasicOp $ SubExp m
      s <- scaleSizesDown lvl elems_per (irregularS ir0)
      o <- scaleSizesDown lvl elems_per (irregularO ir0)
      flags <- subsampleFlags lvl elems_per (irregularF ir0)
      -- The per-segment metadata (outer size @w@): shape and offset arrays.
      bindReshape s_name s (Shape [w])
      bindReshape o_name o (Shape [w])
      -- The per-element flag array (outer size @m@).
      bindReshape f_name flags (Shape [Var m_name])
      -- The concatenated data arrays, reshaped from the flat scalar data to
      -- @[m]b@ (the outer dimension being the freshly bound total size).
      forM_ (zip3 d_names irregs elem_ts) $ \(VName
v, IrregularRep
ir, Type
elem_t) ->
        VName -> VName -> ShapeBase SubExp -> FlattenM ()
forall {m :: * -> *}.
MonadBuilder m =>
VName -> VName -> ShapeBase SubExp -> m ()
bindReshape VName
v (IrregularRep -> VName
irregularD IrregularRep
ir) (Segments -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [VName -> SubExp
Var VName
m_name] ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Segments -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape (Type -> Segments
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
elem_t))
    ([IrregularRep], [Type])
_ -> String -> FlattenM ()
forall a. HasCallStack => String -> a
error String
"flattenFlatMap: FlatMap with no irregular results"
  where
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops
    value_ts :: [ExtType]
value_ts = Int -> [ExtType] -> [ExtType]
forall a. Int -> [a] -> [a]
drop Int
1 (ExtLambda SOACS -> [ExtType]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda SOACS
lam)
    static :: TypeBase (ShapeBase ExtSize) o -> TypeBase (ShapeBase SubExp) o
static =
      TypeBase (ShapeBase SubExp) o
-> Maybe (TypeBase (ShapeBase SubExp) o)
-> TypeBase (ShapeBase SubExp) o
forall a. a -> Maybe a -> a
fromMaybe (String -> TypeBase (ShapeBase SubExp) o
forall a. HasCallStack => String -> a
error String
"flattenFlatMap: existential size.") (Maybe (TypeBase (ShapeBase SubExp) o)
 -> TypeBase (ShapeBase SubExp) o)
-> (TypeBase (ShapeBase ExtSize) o
    -> Maybe (TypeBase (ShapeBase SubExp) o))
-> TypeBase (ShapeBase ExtSize) o
-> TypeBase (ShapeBase SubExp) o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (ShapeBase ExtSize) o
-> Maybe (TypeBase (ShapeBase SubExp) o)
forall o.
TypeBase (ShapeBase ExtSize) o
-> Maybe (TypeBase (ShapeBase SubExp) o)
hasStaticShape
    bindReshape :: VName -> VName -> ShapeBase SubExp -> m ()
bindReshape VName
name VName
arr ShapeBase SubExp
newshape = do
      arr_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
      letBindNames [name] . BasicOp $
        Reshape arr (reshapeAll (arrayShape arr_t) newshape)

-- | Flattening rule for a 'FlatMap' nested inside an enclosing map-nest. A
-- 'FlatMap' is just a nonuniform map with implicit concatenation, so we run it
-- through the ordinary inner-map machinery, which already produces - for a
-- variably-sized result - an 'IrregularRep' whose segment sizes are the
-- per-enclosing-segment concatenated lengths and whose data is the
-- concatenation. That directly gives the data arrays (results 4..) and, as
-- their per-enclosing segment sizes, the total sizes @m@ (result 0).
--
-- The metadata is derived from the shape array (result 1) - the per-iteration
-- output sizes - which the inner-map does not itself surface. The lambda
-- returns exactly that size as its leading result, so mapping it over the nest
-- yields the irregular array of per-iteration sizes, segmented by the 'FlatMap'
-- width. From it we then compute the offset array (result 3), as the
-- per-enclosing exclusive prefix sum, and the flag array (result 2).
--
-- See also Note [FlatMap element counting].
flattenFlatMapNested ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  StmAux () ->
  SubExp ->
  [VName] ->
  ExtLambda SOACS ->
  FlattenM DistEnv
flattenFlatMapNested :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> SubExp
-> [VName]
-> ExtLambda SOACS
-> FlattenM DistEnv
flattenFlatMapNested FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux SubExp
w [VName]
arrs ExtLambda SOACS
lam = do
  size_name <- Name -> FlattenM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"flatmap_sizes"
  data_names <- mapM (const $ newVName "flatmap_data") val_ts
  let size_pe = VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
size_name (Type -> PatElem Type) -> Type -> PatElem Type
forall a b. (a -> b) -> a -> b
$ Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) SubExp
w
      data_pes = (VName -> Type -> PatElem Type)
-> [VName] -> [Type] -> [PatElem Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\VName
v Type
t -> VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
v (Type -> PatElem Type) -> Type -> PatElem Type
forall a b. (a -> b) -> a -> b
$ Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow Type
t SubExp
w) [VName]
data_names [Type]
val_ts
      map_pat = [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ PatElem Type
size_pe PatElem Type -> [PatElem Type] -> [PatElem Type]
forall a. a -> [a] -> [a]
: [PatElem Type]
data_pes
  let lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops
  certifying (distCerts inps aux env) $ do
    reps <- transformMap ops (stmAuxAttrs aux) segments env inps map_pat w arrs map_lam
    case reps of
      ResRep
shape_rep : [ResRep]
data_reps
        | Just ResRep
first_data <- [ResRep] -> Maybe ResRep
forall {b}. [b] -> Maybe b
firstIrregular [ResRep]
data_reps -> do
            -- An 'IrregularRep' always stores primitive data, so for a
            -- non-scalar element type @b@ the data result's segment sizes are
            -- in units of scalars, of which there are 'elems_per' per
            -- b-element. The source 'flatmap' counts b-elements, so we scale
            -- the affected metadata down.
            let elems_per :: TPrimExp Int64 VName
elems_per = ExtLambda SOACS -> TPrimExp Int64 VName
flatMapElemsPer ExtLambda SOACS
lam

            -- The number of b-elements per enclosing segment: the per-segment
            -- totals (result 0), and the segment sizes of the flag array. These
            -- are the data result's segment sizes, scaled to b-element units.
            -- When the data is regular (the fully uniform case), it is the
            -- constant inner size, replicated across the segments.
            s_flag <- case ResRep
first_data of
              Irregular IrregularRep
ir -> SegLevel -> TPrimExp Int64 VName -> VName -> FlattenM VName
scaleSizesDown SegLevel
lvl TPrimExp Int64 VName
elems_per (IrregularRep -> VName
irregularS IrregularRep
ir)
              Regular VName
d -> do
                d_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
d
                let inner = [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ (SubExp -> TPrimExp Int64 VName)
-> Segments -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 (Segments -> [TPrimExp Int64 VName])
-> Segments -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ Int -> Segments -> Segments
forall a. Int -> [a] -> [a]
drop (Segments -> Int
segmentsRank Segments
segments) (Segments -> Segments) -> Segments -> Segments
forall a b. (a -> b) -> a -> b
$ Type -> Segments
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
d_t
                n_elem <- letSubExp "flatmap_n_elem" =<< toExp (inner `div` elems_per)
                letExp "flatmap_n" (BasicOp $ Replicate (segmentsShape segments) n_elem)
            let n_rep = VName -> ResRep
Regular VName
s_flag

            -- The offset array (result 3) has the same structure as the shape
            -- array; its values are the exclusive prefix sum of the shape
            -- within each enclosing segment. We also need the shape as a single
            -- flat array of all per-iteration sizes (for the flags).
            (offset_rep, shape_flat) <- case shape_rep of
              Regular VName
sd -> do
                off <- SegLevel -> Name -> VName -> FlattenM VName
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> Name -> VName -> m VName
genExPrefixSum SegLevel
lvl Name
"flatmap_offset" VName
sd
                sd_t <- lookupType sd
                n_shape <- letSubExp "flatmap_shape_n" =<< toExp (product $ map pe64 $ arrayDims sd_t)
                flat <-
                  letExp "flatmap_shape_flat" . BasicOp $
                    Reshape sd (reshapeAll (arrayShape sd_t) (Shape [n_shape]))
                pure (Regular off, flat)
              Irregular IrregularRep
s_ir -> do
                inc <- SegLevel -> Name -> VName -> VName -> FlattenM VName
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> Name -> VName -> VName -> m VName
genSegPrefixSum SegLevel
lvl Name
"flatmap_offset_inc" (IrregularRep -> VName
irregularF IrregularRep
s_ir) (IrregularRep -> VName
irregularD IrregularRep
s_ir)
                n_off <- arraySize 0 <$> lookupType inc
                off_D <- letExp "flatmap_offset_D" <=< segMap lvl (MkSolo n_off) $ \(MkSolo SubExp
i) -> do
                  a <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"a" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
inc [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
                  b <- letSubExp "b" =<< eIndex (irregularD s_ir) [eSubExp i]
                  off <- letSubExp "off" =<< toExp (pe64 a - pe64 b)
                  pure [subExpRes off]
                pure (Irregular s_ir {irregularD = off_D}, irregularD s_ir)

            -- The flag array (result 2): the segment-start flags over all
            -- b-elements. Every enclosing-segment boundary is also a segment start,
            -- so a single 'doRepIota' over all per-iteration sizes yields the flags
            -- for all enclosing segments at once. Its per-enclosing segment sizes
            -- are @s_flag@ (in b-element units); we build the enclosing flag/offset
            -- structure from those, respecting the invariant that in an
            -- 'IrregularRep' the flag and data arrays share their shape (rather than
            -- borrowing the scalar-unit structure of the data rep).
            (flag_F, flag_O, _) <- doRepIota lvl s_flag
            (flag_D0, _, _) <- doRepIota lvl shape_flat
            -- 'flag_D0' and 'flag_F' both have the total number of b-elements as
            -- their size, but computed by different means; coerce so the rep's flag
            -- and data arrays share a size (the 'IrregularRep' invariant).
            m_flag <- arraySize 0 <$> lookupType flag_F
            flag_D0_t <- lookupType flag_D0
            flag_D <-
              letExp "flatmap_flag_D" . BasicOp $
                Reshape flag_D0 (reshapeAll (arrayShape flag_D0_t) (Shape [m_flag]))
            let flag_rep = IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep) -> IrregularRep -> ResRep
forall a b. (a -> b) -> a -> b
$ VName -> VName -> VName -> VName -> IrregularKind -> IrregularRep
IrregularRep VName
s_flag VName
flag_F VName
flag_O VName
flag_D IrregularKind
Dense

            -- The data arrays (results 4..). The inner-map machinery already
            -- produced, for each, an 'IrregularRep' with primitive
            -- (scalar-unit) data and matching scalar-unit structure arrays.
            -- These are exactly the reps for the source @[m]b@ data arrays, so
            -- we pass them through unchanged; any consumer reconstructs the
            -- @b@-element rows from the primitive rep via the ordinary
            -- irregular-input machinery. See Note [FlatMap element counting].
            let all_reps = ResRep
n_rep ResRep -> [ResRep] -> [ResRep]
forall a. a -> [a] -> [a]
: ResRep
shape_rep ResRep -> [ResRep] -> [ResRep]
forall a. a -> [a] -> [a]
: ResRep
flag_rep ResRep -> [ResRep] -> [ResRep]
forall a. a -> [a] -> [a]
: ResRep
offset_rep ResRep -> [ResRep] -> [ResRep]
forall a. a -> [a] -> [a]
: [ResRep]
data_reps
            insertRepsM (zip (map distResTag res) all_reps) env
      [ResRep]
_ -> String -> FlattenM DistEnv
forall a. HasCallStack => String -> a
error String
"flattenFlatMapNested: FlatMap with no irregular results"
  where
    -- The metadata is derived from a nonuniform result; they all have the same
    -- segment structure, so any of them will do.
    firstIrregular :: [b] -> Maybe b
firstIrregular =
      ((ExtType, b) -> b) -> Maybe (ExtType, b) -> Maybe b
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ExtType, b) -> b
forall a b. (a, b) -> b
snd
        (Maybe (ExtType, b) -> Maybe b)
-> ([b] -> Maybe (ExtType, b)) -> [b] -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ExtType, b) -> Bool) -> [(ExtType, b)] -> Maybe (ExtType, b)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (ExtType -> Bool
flatMapNonuniform (ExtType -> Bool)
-> ((ExtType, b) -> ExtType) -> (ExtType, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ExtType, b) -> ExtType
forall a b. (a, b) -> a
fst)
        ([(ExtType, b)] -> Maybe (ExtType, b))
-> ([b] -> [(ExtType, b)]) -> [b] -> Maybe (ExtType, b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ExtType] -> [b] -> [(ExtType, b)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> [ExtType] -> [ExtType]
forall a. Int -> [a] -> [a]
drop Int
1 (ExtLambda SOACS -> [ExtType]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda SOACS
lam))
    -- This lambda is deliberately ill-formed; see Note [Ill-formed inner-map
    -- lambda].
    map_lam :: Lambda SOACS
map_lam =
      Lambda
        { lambdaParams :: [LParam SOACS]
lambdaParams = ExtLambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams ExtLambda SOACS
lam,
          lambdaReturnType :: [Type]
lambdaReturnType = PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64 Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
val_ts,
          lambdaBody :: Body SOACS
lambdaBody = ExtLambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody ExtLambda SOACS
lam
        }
    -- Only works because we allow just a single Ext in the return type.
    val_ts :: [Type]
val_ts =
      Identity [Type] -> [Type]
forall a. Identity a -> a
runIdentity (Identity [Type] -> [Type])
-> ([ExtType] -> Identity [Type]) -> [ExtType] -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Identity SubExp) -> [ExtType] -> Identity [Type]
forall (m :: * -> *) o.
Monad m =>
(Int -> m SubExp)
-> [TypeBase (ShapeBase ExtSize) o]
-> m [TypeBase (ShapeBase SubExp) o]
instantiateShapes (Identity SubExp -> Int -> Identity SubExp
forall a b. a -> b -> a
const (SubExp -> Identity SubExp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
k)) ([ExtType] -> Identity [Type])
-> ([ExtType] -> [ExtType]) -> [ExtType] -> Identity [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [ExtType] -> [ExtType]
forall a. Int -> [a] -> [a]
drop Int
1 ([ExtType] -> [Type]) -> [ExtType] -> [Type]
forall a b. (a -> b) -> a -> b
$
        ExtLambda SOACS -> [ExtType]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda SOACS
lam
      where
        -- The size of the lambda's irregular results, which it returns before
        -- them, and which is what 'Ext 0' stands for.
        k :: SubExp
k = case Body SOACS -> Result
forall rep res. GBody rep res -> [res]
bodyResult (ExtLambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody ExtLambda SOACS
lam) of
          SubExpRes
r : Result
_ -> SubExpRes -> SubExp
resSubExp SubExpRes
r
          [] -> String -> SubExp
forall a. HasCallStack => String -> a
error String
"flattenFlatMapNested: FlatMap with no results"

-- | Remove certificates that refer to variables free in the lambda (recursing
-- into nested bodies). Such certificates arise on pure operator lambdas (e.g.
-- the combining function of a 'hist' or reduce) through conservative
-- certificate propagation. They are redundant, and cannot be preserved when the
-- operator is lifted into a segmented operation, as the certificate's binding
-- does not survive into the generated kernel.
stripFreeCerts :: Lambda SOACS -> Lambda SOACS
stripFreeCerts :: Lambda SOACS -> Lambda SOACS
stripFreeCerts Lambda SOACS
lam = Lambda SOACS
lam {lambdaBody = onBody (lambdaBody lam)}
  where
    frees :: Names
frees = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
lam
    onBody :: Body rep -> Body rep
onBody Body rep
body = Body rep
body {bodyStms = onStm <$> bodyStms body}
    onStm :: Stm rep -> Stm rep
onStm (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
dec Exp rep
e) = Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec rep)
pat (StmAux (ExpDec rep) -> StmAux (ExpDec rep)
forall {dec}. StmAux dec -> StmAux dec
onDec StmAux (ExpDec rep)
dec) (Exp rep -> Exp rep
onExp Exp rep
e)
    onDec :: StmAux dec -> StmAux dec
onDec StmAux dec
dec =
      StmAux dec
dec {stmAuxCerts = Certs $ filter (`notNameIn` frees) $ unCerts $ stmAuxCerts dec}
    onExp :: Exp rep -> Exp rep
onExp = Identity (Exp rep) -> Exp rep
forall a. Identity a -> a
runIdentity (Identity (Exp rep) -> Exp rep)
-> (Exp rep -> Identity (Exp rep)) -> Exp rep -> Exp rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mapper rep rep Identity -> Exp rep -> Identity (Exp rep)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper rep rep Identity
mapper
    mapper :: Mapper rep rep Identity
mapper = Mapper rep rep Identity
forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper {mapOnBody = const (pure . onBody)}

flattenHist ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  (Pat Type, StmAux ()) ->
  (SubExp, [VName], [Futhark.IR.SOACS.HistOp SOACS], Lambda SOACS) ->
  FlattenM DistEnv
flattenHist :: 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, StmAux ()
aux) (SubExp
w, [VName]
hist_inputs, [HistOp SOACS]
hist_ops0, Lambda SOACS
bucket_fun) = do
  -- The operator is a pure combining function; any certificates on its
  -- statements (from conservative propagation) are redundant and cannot be
  -- preserved when it is lifted into a segmented operation, so drop the ones
  -- referring to variables free in the operator.
  let hist_ops :: [HistOp SOACS]
hist_ops = do
        op <- [HistOp SOACS]
hist_ops0
        pure $
          op
            { Futhark.IR.SOACS.histOp = stripFreeCerts $ Futhark.IR.SOACS.histOp op
            }
  -- TODO: check for suitableUniformOperator.
  let nonuniform :: Bool
nonuniform =
        Bool -> Bool
not ((VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (DistEnv -> DistInputs -> VName -> Bool
isRegularInputArr DistEnv
env DistInputs
inps) [VName]
hist_inputs)
          Bool -> Bool -> Bool
|| DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
w
          Bool -> Bool -> Bool
|| Bool -> Bool
not ((DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
res)
  if Bool
nonuniform
    then do
      gpu_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
      let scope = Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS) -> Scope GPU -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ DistInputs -> Scope GPU
scopeOfDistInputs DistInputs
inps Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope GPU
gpu_scope
      (hist_res, stms) <-
        runBuilderT
          ( auxing aux $
              doHist "nonuniform_hist" hist_ops hist_inputs $ \[LParam (Rep (BuilderT SOACS FlattenM))]
params ->
                (SubExpRes -> SubExp) -> Result -> Segments
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp (Result -> Segments)
-> BuilderT SOACS FlattenM Result
-> BuilderT SOACS FlattenM Segments
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda (Rep (BuilderT SOACS FlattenM))
-> [BuilderT SOACS FlattenM (Exp (Rep (BuilderT SOACS FlattenM)))]
-> BuilderT SOACS FlattenM Result
forall (m :: * -> *).
MonadBuilder m =>
Lambda (Rep m) -> [m (Exp (Rep m))] -> m Result
eLambda Lambda (Rep (BuilderT SOACS FlattenM))
Lambda SOACS
bucket_fun ((Param Type -> BuilderT SOACS FlattenM (Exp SOACS))
-> [Param Type] -> [BuilderT SOACS FlattenM (Exp SOACS)]
forall a b. (a -> b) -> [a] -> [b]
map Param Type
-> BuilderT SOACS FlattenM (Exp (Rep (BuilderT SOACS FlattenM)))
Param Type -> BuilderT SOACS FlattenM (Exp SOACS)
forall (m :: * -> *) t.
MonadBuilder m =>
Param t -> m (Exp (Rep m))
eParam [Param Type]
[LParam (Rep (BuilderT SOACS FlattenM))]
params)
          )
          scope
      let body = Seq (Stm SOACS) -> Result -> Body SOACS
forall res.
IsResult res =>
Seq (Stm SOACS) -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Seq (Stm SOACS)
stms (Result -> Body SOACS) -> Result -> Body SOACS
forall a b. (a -> b) -> a -> b
$ [VName] -> Result
varsRes ([VName] -> Result) -> [VName] -> Result
forall a b. (a -> b) -> a -> b
$ [[VName]] -> [VName]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[VName]]
hist_res
      reps <- distributeAndFlattenBody ops segments "non_uniform_hist_body" env inps res body
      insertRepsM (zip (map distResTag res) reps) env
    else do
      let new_segment :: Segments
new_segment = Segments
segments Segments -> Segments -> Segments
forall a. Semigroup a => a -> a -> a
<> SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w
      lifted_inps <- [VName] -> (VName -> FlattenM VName) -> FlattenM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
hist_inputs ((VName -> FlattenM VName) -> FlattenM [VName])
-> (VName -> FlattenM VName) -> FlattenM [VName]
forall a b. (a -> b) -> a -> b
$ \VName
hist_inp -> do
        t <- DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
hist_inp
        let expectedShape = Segments -> ShapeBase SubExp
segmentsShape Segments
segments ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t
        liftSubExpRegular lvl segments inps env expectedShape (Var hist_inp)
      hist_ops' <- forM hist_ops $ \(Futhark.IR.SOACS.HistOp ShapeBase SubExp
num_bins SubExp
rf [VName]
dests Segments
nes Lambda SOACS
op) -> do
        nes' <- (SubExp -> FlattenM SubExp) -> Segments -> FlattenM Segments
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 (Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM SubExp
readNeutral Segments
segments DistEnv
env DistInputs
inps) Segments
nes
        let rr (DistType Segments
_ Rank
_ Type
t) = Type
t
        let ts = (DistResult -> Type) -> [DistResult] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (DistType -> Type
rr (DistType -> Type)
-> (DistResult -> DistType) -> DistResult -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DistResult -> DistType
distResType) [DistResult]
res
        let expectedShapes = (Type -> ShapeBase SubExp) -> [Type] -> [ShapeBase SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (\Type
t -> Segments -> ShapeBase SubExp
segmentsShape Segments
segments ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t) [Type]
ts
        dests' <- mapM (\(ShapeBase SubExp
shape, VName
var) -> SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> ShapeBase SubExp
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl Segments
segments DistInputs
inps DistEnv
env ShapeBase SubExp
shape (VName -> SubExp
Var VName
var)) (zip expectedShapes dests)
        pure $ Futhark.IR.SOACS.HistOp num_bins rf dests' nes' op
      let free = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
bucket_fun
      let isDest = (VName -> [VName] -> Bool) -> [VName] -> VName -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem ([VName] -> VName -> Bool) -> [VName] -> VName -> Bool
forall a b. (a -> b) -> a -> b
$ (HistOp SOACS -> [VName]) -> [HistOp SOACS] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap HistOp SOACS -> [VName]
forall rep. HistOp rep -> [VName]
Futhark.IR.SOACS.histDest [HistOp SOACS]
hist_ops'
          free_notDest = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Bool
isDest) (Names -> [VName]
namesToList Names
free)
      free_and_sizes <- freeWithTypeDeps inps (namesFromList free_notDest)
      (free_replicated, replicated) <-
        fmap unzip . sequence $
          mapMaybe
            (onMapFreeVarMultiDim lvl segments w env inps)
            free_and_sizes
      (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated
      let readFree Segments
is = Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInputs Segments
new_segment DistEnv
free_env Segments
is DistInputs
free_inputs
      hist_res <-
        certifying (distCerts inps aux env) $
          genUniformSegHist lvl "Uniform_segHist" new_segment hist_ops' (soacsLambdaToGPU bucket_fun) lifted_inps readFree
      pure $ insertRegulars (map distResTag res) hist_res env
  where
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops