-- | Various general utilities used in flattening.
module Futhark.Pass.Flatten.General
  ( -- * Reading inputs
    readInputVar,
    readInputs,
    readInput,
    readNeutral,
    readTypeDims,

    -- * Building blocks
    ensureDenseIrregular,
    liftResult,
    liftDistResultRep,
    liftSubExp,
    liftSubExpPreserveRep,
    liftSubExpRegular,
    liftVarRegular,
    isRegularInputArr,
    liftParam,
    liftRegularParam,
    liftRegResult,
    needsIrregularRetType,
    mkIrregFromReg,
    flattenIrregularRep,
    distCerts,
    dataArr,
    getIrregRep,
    scatterIrregular,
    scatterRegular,
    module Futhark.Pass.Flatten.Monad,
    module Futhark.Pass.Flatten.Builtins,

    -- * Various
    scopeOfDistInputs,
    lookupInputType,
    subExpInputType,
    localiseInputs,
    replicateForDims,
    liftBodyWithDistResults,
    distResultsToResReps,
    resultToResReps,
    resultToResRepsByDistResult,
    irregularRepToFlatArrs,
    distributeAndFlattenBody,
    splitInput,
    isVariant,
    segmentDims,
    flattenDistStm,
    flattenScalarStm,
    distributeBodyWith,
    distributeMapWith,
    atSegLevel,
  )
where

import Control.Monad
import Data.List qualified as L
import Data.Map qualified as M
import Data.Maybe
import Data.Tuple.Solo
import Futhark.IR.GPU
import Futhark.IR.SOACS as SOACS
import Futhark.MonadFreshNames
import Futhark.Pass.Flatten.Builtins
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.Monad
import Futhark.Tools
import Futhark.Transform.Rename (renameExp)
import Futhark.Util (mapAccumLM)
import Futhark.Util.IntegralExp
import Prelude hiding (div, rem)

-- | Write back the irregular results of a branch to a (partially) blank space
-- The `offsets` variable is the offsets of the final result, whereas `irregRep`
-- is the irregular representation of the result.
scatterIrregular ::
  SegLevel ->
  VName ->
  VName ->
  (VName, IrregularRep) ->
  FlattenM VName
scatterIrregular :: SegLevel
-> VName -> VName -> (VName, IrregularRep) -> FlattenM VName
scatterIrregular SegLevel
lvl VName
offsets VName
space (VName
is, IrregularRep
irregRep) = do
  let IrregularRep {irregularS :: IrregularRep -> VName
irregularS = VName
segs, irregularD :: IrregularRep -> VName
irregularD = VName
elems, irregularO :: IrregularRep -> VName
irregularO = VName
off, irregularK :: IrregularRep -> IrregularKind
irregularK = IrregularKind
kind} = IrregularRep
irregRep
  (_, _, ii1) <- SegLevel -> VName -> FlattenM (VName, VName, VName)
doRepIota SegLevel
lvl VName
segs
  (_, _, ii2) <- doSegIota lvl segs
  m <- arraySize 0 <$> lookupType ii1
  letExp "irregular_scatter" <=< genScatter lvl space m $ \SubExp
gtid -> 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
ii1 [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gtid]
    intra_segment <- letSubExp "segment" =<< eIndex ii2 [eSubExp gtid]
    x <- case kind of
      IrregularKind
Dense -> 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
elems [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gtid]
      IrregularKind
Replicated -> do
        o <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"rep_O" (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
off [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
segment]
        letSubExp "x" =<< eIndex elems [toExp $ pe64 o + pe64 intra_segment]
    offset <- letExp "offset" =<< eIndex offsets [eIndex is [eSubExp segment]]
    i <- letExp "i" =<< eBinOp (Add Int64 OverflowUndef) (toExp offset) (eSubExp intra_segment)
    pure (i, x)

-- | Write back the regular results to a (partially) blank space
scatterRegular ::
  SegLevel ->
  VName ->
  (VName, VName) ->
  FlattenM VName
scatterRegular :: SegLevel -> VName -> (VName, VName) -> FlattenM VName
scatterRegular SegLevel
lvl VName
space (VName
is, VName
xs) = do
  dims <- Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (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
xs
  letExp "regular_scatter" <=< genScatterND lvl space dims $ \(SubExp
gtid : [SubExp]
rest) -> 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
xs ((SubExp -> FlattenM (Exp GPU)) -> [SubExp] -> [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 (SubExp
gtid SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
rest))
    i <- letSubExp "i" =<< eIndex is [eSubExp gtid]
    pure (i : rest, x)

ensureDenseIrregular :: SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular :: SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
_ Name
_ rep :: IrregularRep
rep@IrregularRep {irregularK :: IrregularRep -> IrregularKind
irregularK = IrregularKind
Dense} =
  IrregularRep -> FlattenM IrregularRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IrregularRep
rep
ensureDenseIrregular SegLevel
lvl Name
desc rep :: IrregularRep
rep@IrregularRep {} = do
  (new_F, new_O, ii1) <- SegLevel -> VName -> FlattenM (VName, VName, VName)
doRepIota SegLevel
lvl (IrregularRep -> VName
irregularS IrregularRep
rep)
  m <- arraySize 0 <$> lookupType ii1
  new_D <- letExp (desc <> "_dense_D") <=< segMap lvl (MkSolo m) $ \(MkSolo SubExp
i) -> do
    seg <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"seg" (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
ii1 [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    old_off <- letSubExp "old_off" =<< eIndex (irregularO rep) [eSubExp seg]
    new_off <- letSubExp "new_off" =<< eIndex new_O [eSubExp seg]
    j <- letSubExp "j" <=< toExp $ pe64 i - pe64 new_off
    x <- letSubExp "x" =<< eIndex (irregularD rep) [toExp $ pe64 old_off + pe64 j]
    pure [subExpRes x]
  pure $
    IrregularRep
      { irregularS = irregularS rep,
        irregularF = new_F,
        irregularO = new_O,
        irregularD = new_D,
        irregularK = Dense
      }

-- Lift a result of a function.
liftResult :: SegLevel -> Segments -> DistInputs -> DistEnv -> SubExpRes -> FlattenM Result
liftResult :: SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> SubExpRes
-> FlattenM Result
liftResult SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env SubExpRes
res = (VName -> SubExpRes) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (Certs -> SubExp -> SubExpRes
SubExpRes Certs
forall a. Monoid a => a
mempty (SubExp -> SubExpRes) -> (VName -> SubExp) -> VName -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) ([VName] -> Result) -> FlattenM [VName] -> FlattenM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FlattenM [VName]
vs
  where
    vs :: FlattenM [VName]
vs = do
      (_, rep) <- SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> SubExp
-> FlattenM (Type, ResRep)
liftSubExp SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env (SubExpRes -> SubExp
resSubExp SubExpRes
res)
      case rep of
        Regular VName
v -> [VName] -> FlattenM [VName]
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName
v]
        Irregular IrregularRep
irreg -> IrregularRep -> FlattenM [VName]
forall {m :: * -> *}. MonadBuilder m => IrregularRep -> m [VName]
mkIrrep IrregularRep
irreg
    mkIrrep :: IrregularRep -> m [VName]
mkIrrep
      ( IrregularRep
          { irregularS :: IrregularRep -> VName
irregularS = VName
segs,
            irregularF :: IrregularRep -> VName
irregularF = VName
flags,
            irregularO :: IrregularRep -> VName
irregularO = VName
offsets,
            irregularD :: IrregularRep -> VName
irregularD = VName
elems
          }
        ) = do
        flags_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
flags
        t <- lookupType elems
        num_data <- letExp "num_data" =<< toExp (product $ map pe64 $ arrayDims t)
        let shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [VName -> SubExp
Var VName
num_data]
        flags' <- letExp "flags" $ BasicOp $ Reshape flags $ reshapeAll (arrayShape flags_t) shape
        elems' <- letExp "elems" $ BasicOp $ Reshape elems $ reshapeAll (arrayShape t) shape
        pure [num_data, segs, flags', offsets, elems']

needsIrregularRetType :: DistInputs -> RetType SOACS -> Bool
needsIrregularRetType :: DistInputs -> RetType SOACS -> Bool
needsIrregularRetType DistInputs
inps = (ExtSize -> Bool) -> [ExtSize] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ExtSize -> Bool
needsIrregularDim ([ExtSize] -> Bool) -> (ExtType -> [ExtSize]) -> ExtType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtType -> [ExtSize]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims
  where
    needsIrregularDim :: ExtSize -> Bool
needsIrregularDim Ext {} = Bool
True
    needsIrregularDim (Free SubExp
se) = DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
se

liftRegResult :: SegLevel -> Segments -> SubExp -> DistInputs -> DistEnv -> RetType SOACS -> SubExpRes -> FlattenM Result
liftRegResult :: SegLevel
-> [SubExp]
-> SubExp
-> DistInputs
-> DistEnv
-> RetType SOACS
-> SubExpRes
-> FlattenM Result
liftRegResult SegLevel
lvl [SubExp]
segments SubExp
num_segments DistInputs
inps DistEnv
env RetType SOACS
rettype SubExpRes
res
  | DistInputs -> RetType SOACS -> Bool
needsIrregularRetType DistInputs
inps RetType SOACS
rettype = case SubExpRes -> SubExp
resSubExp SubExpRes
res of
      Var VName
v -> do
        irreg <- SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
v
        varsRes <$> irregularRepToFlatArrs num_segments irreg
      Constant {} ->
        String -> FlattenM Result
forall a. HasCallStack => String -> a
error String
"liftRegResult: irregular result is not a variable"
  | Bool
otherwise = do
      let res_se :: SubExp
res_se = SubExpRes -> SubExp
resSubExp SubExpRes
res
      res_t <- DistInputs -> SubExp -> FlattenM Type
subExpInputType DistInputs
inps SubExp
res_se
      let expectedShape = [SubExp] -> Shape
segmentsShape [SubExp]
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
res_t
      lifted_res <- liftSubExpRegular lvl segments inps env expectedShape res_se
      pure [SubExpRes mempty (Var lifted_res)]

mkIrregFromReg ::
  SegLevel ->
  Segments ->
  VName ->
  FlattenM IrregularRep
mkIrregFromReg :: SegLevel -> [SubExp] -> VName -> FlattenM IrregularRep
mkIrregFromReg SegLevel
lvl [SubExp]
segments VName
arr = do
  arr_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
  num_segments <-
    letSubExp "reg_num_segments" <=< toExp $ product $ segmentDims segments
  segment_size <-
    letSubExp "reg_seg_size" <=< toExp . product . map pe64 $
      drop (segmentsRank segments) (arrayDims arr_t)
  arr_S <-
    letExp "reg_segments" . BasicOp $
      Replicate (Shape [num_segments]) segment_size
  num_elems <-
    letSubExp "reg_num_elems" <=< toExp $ product $ map pe64 $ arrayDims arr_t
  arr_D <-
    letExp "reg_D" . BasicOp $
      Reshape arr (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_O <- letExp "reg_O" <=< segMap lvl (MkSolo num_segments) $ \(MkSolo SubExp
i) -> do
    offset <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"offset" (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 a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
segment_size
    pure [subExpRes offset]
  pure $
    IrregularRep
      { irregularS = arr_S,
        irregularF = arr_F,
        irregularO = arr_O,
        irregularD = arr_D,
        irregularK = Dense
      }

readIrregularInput ::
  Segments ->
  [SubExp] ->
  VName ->
  Type ->
  IrregularRep ->
  FlattenM VName
readIrregularInput :: [SubExp]
-> [SubExp] -> VName -> Type -> IrregularRep -> FlattenM VName
readIrregularInput [SubExp]
segments [SubExp]
is VName
v Type
t (IrregularRep VName
_ VName
_ VName
v_O VName
v_D IrregularKind
_) = do
  offset <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"offset" (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
v_O [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] -> [SubExp] -> TPrimExp Int64 VName
flatSegmentIndex [SubExp]
segments [SubExp]
is]
  case arrayDims t of
    [] -> 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
"_inp") (Exp GPU -> FlattenM VName) -> FlattenM (Exp GPU) -> FlattenM VName
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
v_D [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
offset]
    [SubExp
num_elems] -> do
      let slice :: Slice SubExp
slice = [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice [SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice SubExp
offset SubExp
num_elems (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)]
      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
"_inp") (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 -> Slice SubExp -> BasicOp
Index VName
v_D Slice SubExp
slice
    [SubExp]
_ -> do
      num_elems <-
        Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"num_elems" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp ([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)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 ([SubExp] -> [TPrimExp Int64 VName])
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
t)
      let slice = [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice [SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice SubExp
offset SubExp
num_elems (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)]
      v_flat <-
        letExp (baseName v <> "_flat") $ BasicOp $ Index v_D slice
      v_flat_t <- lookupType v_flat
      letExp (baseName v <> "_inp") . BasicOp $
        Reshape v_flat (reshapeAll (arrayShape v_flat_t) (arrayShape t))

readInputVar :: Segments -> DistEnv -> [SubExp] -> DistInputs -> VName -> FlattenM VName
readInputVar :: [SubExp]
-> DistEnv -> [SubExp] -> DistInputs -> VName -> FlattenM VName
readInputVar [SubExp]
segments DistEnv
env [SubExp]
is DistInputs
inputs VName
v =
  case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inputs of
    Maybe DistInput
Nothing -> VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
    Just (DistInputFree VName
arr Type
t)
      | Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t -> VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
arr
      | Bool
otherwise -> 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) -> FlattenM (Exp GPU) -> FlattenM VName
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)) -> [SubExp] -> [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 [SubExp]
is)
    Just (DistInput ResTag
rt Type
t) -> do
      case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
        Regular VName
arr
          | Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t -> VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
arr
          | Bool
otherwise -> 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) -> FlattenM (Exp GPU) -> FlattenM VName
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)) -> [SubExp] -> [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 [SubExp]
is)
        Irregular IrregularRep
irreg -> [SubExp]
-> [SubExp] -> VName -> Type -> IrregularRep -> FlattenM VName
readIrregularInput [SubExp]
segments [SubExp]
is VName
v Type
t IrregularRep
irreg

readInput :: Segments -> DistEnv -> [SubExp] -> DistInputs -> SubExp -> FlattenM SubExp
readInput :: [SubExp]
-> DistEnv -> [SubExp] -> DistInputs -> SubExp -> FlattenM SubExp
readInput [SubExp]
_ DistEnv
_ [SubExp]
_ DistInputs
_ (Constant PrimValue
x) =
  SubExp -> FlattenM SubExp
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> FlattenM SubExp) -> SubExp -> FlattenM SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
x
readInput [SubExp]
segments DistEnv
env [SubExp]
is DistInputs
inputs (Var VName
v) =
  VName -> SubExp
Var (VName -> SubExp) -> FlattenM VName -> FlattenM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExp]
-> DistEnv -> [SubExp] -> DistInputs -> VName -> FlattenM VName
readInputVar [SubExp]
segments DistEnv
env [SubExp]
is DistInputs
inputs VName
v

-- | Read the neutral element of a reduction or scan operator. The neutral
-- element of an operator is unique because we assume uniform operators, so a
-- valid program must have the same value in every segment, and we can read it
-- from the first one. This is the case even if the neutral element is
-- nonuniform, which is useless but allowed. The segment space may however be
-- empty, in which case we produce a blank value instead; it is never used then,
-- but an unconditional read would be out of bounds.
readNeutral :: Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM SubExp
readNeutral :: [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM SubExp
readNeutral [SubExp]
segments DistEnv
env DistInputs
inps SubExp
ne
  | Var VName
v <- SubExp
ne,
    Just DistInput
_ <- VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps = do
      ne_t <- DistInputs -> SubExp -> FlattenM Type
subExpInputType DistInputs
inps SubExp
ne
      n <- letSubExp "num_segments" =<< toExp (segmentCount segments)
      letSubExp (baseName v <> "_ne")
        =<< eIf
          (toExp $ pe64 n .==. 0)
          (eBody [eBlank ne_t])
          (eBody [eSubExp =<< readInput segments env zeros inps ne])
  | Bool
otherwise =
      [SubExp]
-> DistEnv -> [SubExp] -> DistInputs -> SubExp -> FlattenM SubExp
readInput [SubExp]
segments DistEnv
env [SubExp]
zeros DistInputs
inps SubExp
ne
  where
    zeros :: [SubExp]
zeros = Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate ([SubExp] -> Int
segmentsRank [SubExp]
segments) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0)

readTypeDims ::
  Segments ->
  DistEnv ->
  [SubExp] ->
  DistInputs ->
  TypeBase Shape o ->
  FlattenM [SubExp]
readTypeDims :: forall o.
[SubExp]
-> DistEnv
-> [SubExp]
-> DistInputs
-> TypeBase Shape o
-> FlattenM [SubExp]
readTypeDims [SubExp]
segments DistEnv
env [SubExp]
is DistInputs
inputs =
  (SubExp -> FlattenM SubExp) -> [SubExp] -> FlattenM [SubExp]
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 ([SubExp]
-> DistEnv -> [SubExp] -> DistInputs -> SubExp -> FlattenM SubExp
readInput [SubExp]
segments DistEnv
env [SubExp]
is DistInputs
inputs) ([SubExp] -> FlattenM [SubExp])
-> (TypeBase Shape o -> [SubExp])
-> TypeBase Shape o
-> FlattenM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape o -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims

segmentDims :: Segments -> [TPrimExp Int64 VName]
segmentDims :: [SubExp] -> [TPrimExp Int64 VName]
segmentDims = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 ([SubExp] -> [TPrimExp Int64 VName])
-> ([SubExp] -> [SubExp]) -> [SubExp] -> [TPrimExp Int64 VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (Shape -> [SubExp]) -> ([SubExp] -> Shape) -> [SubExp] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SubExp] -> Shape
segmentsShape

flatSegmentIndex :: Segments -> [SubExp] -> TPrimExp Int64 VName
flatSegmentIndex :: [SubExp] -> [SubExp] -> TPrimExp Int64 VName
flatSegmentIndex [SubExp]
segments = [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall num. IntegralExp num => [num] -> [num] -> num
flattenIndex ([SubExp] -> [TPrimExp Int64 VName]
segmentDims [SubExp]
segments) ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
-> ([SubExp] -> [TPrimExp Int64 VName])
-> [SubExp]
-> TPrimExp Int64 VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64

readInputs :: Segments -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()
readInputs :: [SubExp] -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()
readInputs [SubExp]
segments DistEnv
env [SubExp]
is = ((VName, DistInput) -> FlattenM ()) -> DistInputs -> FlattenM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (VName, DistInput) -> FlattenM ()
onInput
  where
    bindInputName :: VName -> Exp (Rep m) -> m ()
bindInputName VName
v Exp (Rep m)
e
      | VName
v VName -> Names -> Bool
`nameIn` Exp (Rep m) -> Names
forall a. FreeIn a => a -> Names
freeIn Exp (Rep m)
e = do
          v' <- Name -> Exp (Rep m) -> m 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
"_inp") Exp (Rep m)
e
          letBindNames [v] $ BasicOp $ SubExp $ Var v'
      | Bool
otherwise =
          [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] Exp (Rep m)
e
    onInput :: (VName, DistInput) -> FlattenM ()
onInput (VName
v, DistInputFree VName
arr Type
t) =
      VName -> Exp (Rep FlattenM) -> FlattenM ()
forall {m :: * -> *}.
MonadBuilder m =>
VName -> Exp (Rep m) -> m ()
bindInputName VName
v
        (Exp GPU -> FlattenM ()) -> FlattenM (Exp GPU) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
          then SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (VName -> SubExp
Var VName
arr)
          else 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)) -> [SubExp] -> [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 [SubExp]
is)
    onInput (VName
v, DistInput ResTag
rt Type
t) =
      case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
        Regular VName
arr ->
          VName -> Exp (Rep FlattenM) -> FlattenM ()
forall {m :: * -> *}.
MonadBuilder m =>
VName -> Exp (Rep m) -> m ()
bindInputName VName
v
            (Exp (Rep FlattenM) -> FlattenM ())
-> FlattenM (Exp (Rep FlattenM)) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
              then 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
              else 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)) -> [SubExp] -> [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 [SubExp]
is)
        Irregular IrregularRep
irreg ->
          [SubExp]
-> [SubExp] -> VName -> Type -> IrregularRep -> FlattenM VName
readIrregularInput [SubExp]
segments [SubExp]
is VName
v Type
t IrregularRep
irreg FlattenM VName
-> (VName -> FlattenM (Exp GPU)) -> FlattenM (Exp GPU)
forall a b. FlattenM a -> (a -> FlattenM b) -> FlattenM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SubExp -> FlattenM (Exp (Rep FlattenM))
SubExp -> FlattenM (Exp GPU)
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp -> FlattenM (Exp GPU))
-> (VName -> SubExp) -> VName -> FlattenM (Exp GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var FlattenM (Exp GPU) -> (Exp GPU -> FlattenM ()) -> FlattenM ()
forall a b. FlattenM a -> (a -> FlattenM b) -> FlattenM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= VName -> Exp (Rep FlattenM) -> FlattenM ()
forall {m :: * -> *}.
MonadBuilder m =>
VName -> Exp (Rep m) -> m ()
bindInputName VName
v

scopeOfDistInputs :: DistInputs -> Scope GPU
scopeOfDistInputs :: DistInputs -> Scope GPU
scopeOfDistInputs = [Param Type] -> Scope GPU
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param Type] -> Scope GPU)
-> (DistInputs -> [Param Type]) -> DistInputs -> Scope GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, DistInput) -> Param Type) -> DistInputs -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (VName, DistInput) -> Param Type
f
  where
    f :: (VName, DistInput) -> Param Type
f (VName
v, DistInput
inp) = Attrs -> VName -> Type -> Param Type
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
v (DistInput -> Type
distInputType DistInput
inp)

lookupInputType :: DistInputs -> VName -> FlattenM Type
lookupInputType :: DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
v =
  case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps of
    Just DistInput
inp -> Type -> FlattenM Type
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> FlattenM Type) -> Type -> FlattenM Type
forall a b. (a -> b) -> a -> b
$ DistInput -> Type
distInputType DistInput
inp
    Maybe DistInput
Nothing -> VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v

subExpInputType :: DistInputs -> SubExp -> FlattenM Type
subExpInputType :: DistInputs -> SubExp -> FlattenM Type
subExpInputType DistInputs
_ (Constant PrimValue
val) =
  Type -> FlattenM Type
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> FlattenM Type) -> Type -> FlattenM Type
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
val
subExpInputType DistInputs
inps (Var VName
v) =
  DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
v

isVariant :: DistInputs -> SubExp -> Bool
isVariant :: DistInputs -> SubExp -> Bool
isVariant DistInputs
_ (Constant PrimValue
_) = Bool
False
isVariant DistInputs
inps (Var VName
v) = Maybe DistInput -> Bool
forall a. Maybe a -> Bool
isJust (Maybe DistInput -> Bool) -> Maybe DistInput -> Bool
forall a b. (a -> b) -> a -> b
$ VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps

liftDistResultRep ::
  SegLevel ->
  Segments ->
  DistInputs ->
  DistEnv ->
  DistResult ->
  SubExpRes ->
  FlattenM ResRep
liftDistResultRep :: SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> DistResult
-> SubExpRes
-> FlattenM ResRep
liftDistResultRep SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env DistResult
dist_res SubExpRes
res
  | DistResult -> Bool
isRegularDistResult DistResult
dist_res = do
      let (DistType [SubExp]
_ Rank
_ Type
t) = DistResult -> DistType
distResType DistResult
dist_res
          expectedShape :: Shape
expectedShape = [SubExp] -> Shape
segmentsShape [SubExp]
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t
      VName -> ResRep
Regular (VName -> ResRep) -> FlattenM VName -> FlattenM ResRep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
expectedShape (SubExpRes -> SubExp
resSubExp SubExpRes
res)
  | Bool
otherwise =
      case SubExpRes -> SubExp
resSubExp SubExpRes
res of
        Var VName
v -> do
          rep <- SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
v
          pure $ Irregular rep
        SubExp
_ -> String -> FlattenM ResRep
forall a. HasCallStack => String -> a
error String
"liftBranchResultRep: irregular result is not a variable"

liftDistResult :: SegLevel -> Segments -> DistInputs -> DistEnv -> DistResult -> SubExpRes -> FlattenM Result
liftDistResult :: SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> DistResult
-> SubExpRes
-> FlattenM Result
liftDistResult SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env DistResult
dist_res SubExpRes
res =
  if DistResult -> Bool
isRegularDistResult DistResult
dist_res
    then do
      let (DistType [SubExp]
_ Rank
_ Type
t) = DistResult -> DistType
distResType DistResult
dist_res
      let expectedShape :: Shape
expectedShape = [SubExp] -> Shape
segmentsShape [SubExp]
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t
      v <- SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
expectedShape (SubExpRes -> SubExp
resSubExp SubExpRes
res)
      pure [SubExpRes mempty (Var v)]
    else case SubExpRes -> SubExp
resSubExp SubExpRes
res of
      Var VName
v -> do
        irreg <- SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
v
        pure $ map (SubExpRes mempty . Var) [irregularS irreg, irregularF irreg, irregularO irreg, irregularD irreg]
      SubExp
_ -> String -> FlattenM Result
forall a. HasCallStack => String -> a
error String
"liftDistResult: irregular result is not a variable"

liftBodyWithDistResults :: FlattenOps -> Segments -> DistInputs -> DistEnv -> DistStms -> [DistResult] -> Result -> FlattenM Result
liftBodyWithDistResults :: FlattenOps
-> [SubExp]
-> DistInputs
-> DistEnv
-> DistStms
-> [DistResult]
-> Result
-> FlattenM Result
liftBodyWithDistResults FlattenOps
ops [SubExp]
segments DistInputs
inputs DistEnv
env DistStms
dstms [DistResult]
dist_res Result
result = do
  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 -> [SubExp] -> DistEnv -> DistStm -> FlattenM DistEnv
flattenDistStm FlattenOps
ops [SubExp]
segments) DistEnv
env DistStms
dstms
  result' <- zipWithM (liftDistResult (flattenSegLevel ops) segments inputs env') dist_res result
  pure $ concat result'

distResultsToResReps :: [DistResult] -> [VName] -> [ResRep]
distResultsToResReps :: [DistResult] -> [VName] -> [ResRep]
distResultsToResReps [DistResult]
dist_res [VName]
results =
  ([VName], [ResRep]) -> [ResRep]
forall a b. (a, b) -> b
snd (([VName], [ResRep]) -> [ResRep])
-> ([VName], [ResRep]) -> [ResRep]
forall a b. (a -> b) -> a -> b
$ ([VName] -> DistResult -> ([VName], ResRep))
-> [VName] -> [DistResult] -> ([VName], [ResRep])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL [VName] -> DistResult -> ([VName], ResRep)
f [VName]
results [DistResult]
dist_res
  where
    f :: [VName] -> DistResult -> ([VName], ResRep)
f [VName]
rs DistResult
dist_res' =
      if DistResult -> Bool
isRegularDistResult DistResult
dist_res'
        then
          let (VName
v : [VName]
rs') = [VName]
rs
           in ([VName]
rs', VName -> ResRep
Regular VName
v)
        else
          let (VName
segs : VName
flags : VName
offsets : VName
elems : [VName]
rs') = [VName]
rs
           in ([VName]
rs', IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep) -> IrregularRep -> ResRep
forall a b. (a -> b) -> a -> b
$ VName -> VName -> VName -> VName -> IrregularKind -> IrregularRep
IrregularRep VName
segs VName
flags VName
offsets VName
elems IrregularKind
Dense)

-- | Convert an irregular representation to its flat constituents (number of
-- data elements, segments, flags, offsets, elements), with the structure arrays
-- reshaped to be one-dimensional of the given width. This is the form in which
-- irregular values are passed to lifted functions and carried through loops.
irregularRepToFlatArrs :: SubExp -> IrregularRep -> FlattenM [VName]
irregularRepToFlatArrs :: SubExp -> IrregularRep -> FlattenM [VName]
irregularRepToFlatArrs SubExp
w (IrregularRep VName
segs VName
flags VName
offsets VName
elems IrregularKind
_) = do
  t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
elems
  t_o <- lookupType offsets
  flags_t <- lookupType flags
  num_data <- letExp "num_data" =<< toExp (product $ map pe64 $ arrayDims t)
  let shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [VName -> SubExp
Var VName
num_data]
  flags' <- letExp "flags" $ BasicOp $ Reshape flags $ reshapeAll (arrayShape flags_t) shape
  elems' <- letExp "elems" $ BasicOp $ Reshape elems $ reshapeAll (arrayShape t) shape
  segs' <- letExp "segs" $ BasicOp $ Reshape segs $ reshapeAll (arrayShape t_o) (Shape [w])
  offsets' <- letExp "offsets" $ BasicOp $ Reshape offsets $ reshapeAll (arrayShape t_o) (Shape [w])
  pure [num_data, segs', flags', offsets', elems']

-- | Distribute a body and lift the distributed statements, giving
-- back representations of the body results.
distributeAndFlattenBody ::
  FlattenOps ->
  Segments ->
  Name ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Body SOACS ->
  FlattenM [ResRep]
distributeAndFlattenBody :: FlattenOps
-> [SubExp]
-> Name
-> DistEnv
-> DistInputs
-> [DistResult]
-> Body SOACS
-> FlattenM [ResRep]
distributeAndFlattenBody FlattenOps
ops [SubExp]
segments Name
desc DistEnv
env DistInputs
inps [DistResult]
res Body SOACS
body = do
  scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  (inps_local, env_local, _) <- localiseInputs env inps
  let (inps_dist, dstms) = distributeBodyWith ops scope segments inps_local body
  lifted_res <- liftBodyWithDistResults ops segments inps_dist env_local dstms res (bodyResult body)
  lifted_vs <- mapM (letExp desc <=< toExp . resSubExp) lifted_res
  pure $ distResultsToResReps res lifted_vs

-- | Take the elements at index @is@ from an input @v@. The representation of
-- @v@ can be overridden through the provided mapping, which takes precedence
-- over what the environment says.
splitInput ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  VName ->
  M.Map VName ResRep ->
  VName ->
  FlattenM (Type, VName, ResRep)
splitInput :: SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> Map VName ResRep
-> VName
-> FlattenM (Type, VName, ResRep)
splitInput SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
is Map VName ResRep
acc_reps VName
v = do
  (t, rep0) <- [SubExp]
-> DistInputs -> DistEnv -> SubExp -> FlattenM (Type, ResRep)
liftSubExpPreserveRep [SubExp]
segments DistInputs
inps DistEnv
env (VName -> SubExp
Var VName
v)
  let rep = ResRep -> VName -> Map VName ResRep -> ResRep
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault ResRep
rep0 VName
v Map VName ResRep
acc_reps
  (t,v,) <$> case rep of
    Regular VName
arr -> do
      if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
        then
          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
        else do
          -- In the regular case we just take the elements
          -- of the array given by `is`
          n <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"n" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (SubExp -> FlattenM (Exp (Rep FlattenM))
SubExp -> FlattenM (Exp GPU)
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp (SubExp -> FlattenM (Exp GPU))
-> (Type -> SubExp) -> Type -> FlattenM (Exp GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 (Type -> FlattenM (Exp GPU)) -> FlattenM Type -> FlattenM (Exp GPU)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
is)
          inner_dims <- drop (segmentsRank segments) . arrayDims <$> lookupType arr
          -- Do the segMap over all dims, so the inner dimensions
          -- are gathered in parallel
          arr' <- letExp "split_arr" <=< segMap lvl (n : inner_dims) $ \(SubExp
i : [SubExp]
js) -> do
            idx <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"idx" (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
is [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
            let arr_is = [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex ([SubExp] -> [TPrimExp Int64 VName]
segmentDims [SubExp]
segments) (SubExp -> TPrimExp Int64 VName
pe64 SubExp
idx)
            subExpsRes . pure <$> (letSubExp "arr" =<< eIndex arr (map toExp arr_is ++ map eSubExp js))
          pure $ Regular arr'
    Irregular (IrregularRep VName
segs VName
flags VName
offsets VName
elems IrregularKind
_) -> do
      -- In the irregular case we take the elements
      -- of the `segs` array given by `is` like in the regular case
      n <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"n" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (SubExp -> FlattenM (Exp (Rep FlattenM))
SubExp -> FlattenM (Exp GPU)
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp (SubExp -> FlattenM (Exp GPU))
-> (Type -> SubExp) -> Type -> FlattenM (Exp GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 (Type -> FlattenM (Exp GPU)) -> FlattenM Type -> FlattenM (Exp GPU)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
is)
      segs' <- letExp "split_segs" <=< segMap lvl (MkSolo n) $ \(MkSolo SubExp
i) -> do
        idx <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"idx" (Exp GPU -> FlattenM VName) -> FlattenM (Exp GPU) -> FlattenM VName
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
is [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
        subExpsRes . pure <$> (letSubExp "segs" =<< eIndex segs [toExp idx])
      -- From this we calculate the offsets and number of elements
      (_, offsets', num_data) <- exScanAndSum lvl segs'
      (_, _, ii1) <- doRepIota lvl segs'
      (_, _, ii2) <- doSegIota lvl segs'
      -- We then take the elements we need from `elems` and `flags`
      -- For each index `i`, we roughly:
      -- Get the offset of the segment we want to copy by indexing
      -- `offsets` through `is` further through `ii1` i.e.
      -- `offset = offsets[is[ii1[i]]]`
      -- We then add `ii2[i]` to `offset`
      -- and use that to index into `elems` and `flags`.
      ~[flags', elems'] <- letTupExp "split_F_data" <=< segMap lvl (MkSolo num_data) $ \(MkSolo SubExp
i) -> do
        offset <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"offset" (Exp GPU -> FlattenM VName) -> FlattenM (Exp GPU) -> FlattenM VName
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
offsets [VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
is [VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ii1 [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]]]
        idx <- letExp "idx" =<< eBinOp (Add Int64 OverflowUndef) (toExp offset) (eIndex ii2 [eSubExp i])
        flags_split <- letSubExp "flags" =<< eIndex flags [toExp idx]
        elems_split <- letSubExp "elems" =<< eIndex elems [toExp idx]
        pure $ subExpsRes [flags_split, elems_split]
      pure $
        Irregular $
          IrregularRep
            { irregularS = segs',
              irregularF = flags',
              irregularO = offsets',
              irregularD = elems',
              irregularK = Dense
            }

-- | Flatten the arrays of an IrregularRep to be entirely one-dimensional.
flattenIrregularRep :: SegLevel -> IrregularRep -> FlattenM IrregularRep
flattenIrregularRep :: SegLevel -> IrregularRep -> FlattenM IrregularRep
flattenIrregularRep SegLevel
lvl ir :: IrregularRep
ir@(IrregularRep VName
shape VName
_ VName
offsets VName
elems IrregularKind
kind) = do
  elems_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
elems
  if arrayRank elems_t == 1
    then pure ir
    else do
      n <- arraySize 0 <$> lookupType shape
      m' <- letSubExp "flat_m" <=< toExp $ product $ map pe64 $ arrayDims elems_t
      elems' <-
        letExp (baseName elems <> "_flat") . BasicOp $
          Reshape elems (reshapeAll (arrayShape elems_t) (Shape [m']))
      let inner_size = [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)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 ([SubExp] -> [TPrimExp Int64 VName])
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [SubExp]
forall a. HasCallStack => [a] -> [a]
tail ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
elems_t
      ~[shape', offsets'] <-
        letTupExp (baseName shape <> "_flat_metadata")
          <=< renameExp
          <=< segMap lvl (MkSolo n)
          $ \(MkSolo SubExp
i) -> do
            old_shape <-
              Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"old_shape" (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
shape [SubExp -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp SubExp
i]
            old_offset <-
              letSubExp "old_offset" =<< eIndex offsets [toExp i]

            segment_shape <-
              letSubExp "segment_shape" <=< toExp $ (pe64 old_shape * inner_size)
            segment_offset <-
              letSubExp "segment_offset" <=< toExp $ (pe64 old_offset * inner_size)
            pure $ subExpsRes [segment_shape, segment_offset]
      flags' <- genFlags lvl m' offsets'
      pure $ IrregularRep shape' flags' offsets' elems' kind

-- If the sub-expression is a constant, replicate it to match the shape of `segments`
-- If it's a variable, lookup the variable in the dist inputs and dist env,
-- and if it can't be found it is a free variable, so we replicate it to match the shape of `segments`.
liftSubExp :: SegLevel -> Segments -> DistInputs -> DistEnv -> SubExp -> FlattenM (Type, ResRep)
liftSubExp :: SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> SubExp
-> FlattenM (Type, ResRep)
liftSubExp SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env SubExp
se = case SubExp
se of
  c :: SubExp
c@(Constant PrimValue
prim) ->
    let t :: TypeBase shape o
t = PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim (PrimType -> TypeBase shape o) -> PrimType -> TypeBase shape o
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
prim
     in ((Type
forall {shape} {o}. TypeBase shape o
t,) (ResRep -> (Type, ResRep))
-> (VName -> ResRep) -> VName -> (Type, ResRep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> ResRep
Regular (VName -> (Type, ResRep))
-> FlattenM VName -> FlattenM (Type, ResRep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"lifted_const" (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
$ Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
segmentsShape [SubExp]
segments) SubExp
c))
  Var VName
v -> case VName -> Map VName (Type, ResRep) -> Maybe (Type, ResRep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName (Type, ResRep) -> Maybe (Type, ResRep))
-> Map VName (Type, ResRep) -> Maybe (Type, ResRep)
forall a b. (a -> b) -> a -> b
$ DistInputs -> DistEnv -> Map VName (Type, ResRep)
inputReps DistInputs
inps DistEnv
env of
    Just (Type
t, Regular VName
v') -> do
      (Type
t,)
        (ResRep -> (Type, ResRep))
-> FlattenM ResRep -> FlattenM (Type, ResRep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Type
t of
          Prim {} -> 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'
          Array {} -> IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep)
-> FlattenM IrregularRep -> FlattenM ResRep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel -> [SubExp] -> VName -> FlattenM IrregularRep
mkIrregFromReg SegLevel
lvl [SubExp]
segments VName
v'
          Acc {} -> 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'
          Mem {} -> String -> FlattenM ResRep
forall a. HasCallStack => String -> a
error String
"liftSubExp: Mem"
    Just (Type
t, Irregular IrregularRep
irreg) -> do
      irreg' <- SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl Name
"lifted_irreg" IrregularRep
irreg
      (t,)
        <$> case t of
          Prim {} -> 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 -> ResRep) -> VName -> ResRep
forall a b. (a -> b) -> a -> b
$ IrregularRep -> VName
irregularD IrregularRep
irreg'
          Array {} -> 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
$ IrregularRep -> ResRep
Irregular IrregularRep
irreg'
          Acc {} -> String -> FlattenM ResRep
forall a. HasCallStack => String -> a
error String
"liftSubExp: Irregular Acc"
          Mem {} -> String -> FlattenM ResRep
forall a. HasCallStack => String -> a
error String
"liftSubExp: Mem"
    Maybe (Type, ResRep)
Nothing -> do
      t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
      v' <- letExp "free_replicated" $ BasicOp $ Replicate (segmentsShape segments) (Var v)
      (t,)
        <$> case t of
          Prim {} -> 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'
          Array {} -> IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep)
-> FlattenM IrregularRep -> FlattenM ResRep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel -> [SubExp] -> VName -> FlattenM IrregularRep
mkIrregFromReg SegLevel
lvl [SubExp]
segments VName
v'
          Acc {} -> 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'
          Mem {} -> String -> FlattenM ResRep
forall a. HasCallStack => String -> a
error String
"getRepSubExp: Mem"

liftSubExpPreserveRep :: Segments -> DistInputs -> DistEnv -> SubExp -> FlattenM (Type, ResRep)
liftSubExpPreserveRep :: [SubExp]
-> DistInputs -> DistEnv -> SubExp -> FlattenM (Type, ResRep)
liftSubExpPreserveRep [SubExp]
segments DistInputs
inps DistEnv
env SubExp
se = case SubExp
se of
  c :: SubExp
c@(Constant PrimValue
prim) ->
    let t :: TypeBase shape o
t = PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim (PrimType -> TypeBase shape o) -> PrimType -> TypeBase shape o
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
prim
     in do
          v <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"lifted_const" (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
$ Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
segmentsShape [SubExp]
segments) SubExp
c
          pure (t, Regular v)
  Var VName
v -> case VName -> Map VName (Type, ResRep) -> Maybe (Type, ResRep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName (Type, ResRep) -> Maybe (Type, ResRep))
-> Map VName (Type, ResRep) -> Maybe (Type, ResRep)
forall a b. (a -> b) -> a -> b
$ DistInputs -> DistEnv -> Map VName (Type, ResRep)
inputReps DistInputs
inps DistEnv
env of
    Just (Type
t, ResRep
rep) -> (Type, ResRep) -> FlattenM (Type, ResRep)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
t, ResRep
rep)
    Maybe (Type, ResRep)
Nothing -> do
      t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
      v' <- letExp "free_replicated" $ BasicOp $ Replicate (segmentsShape segments) (Var v)
      pure (t, Regular v')

-- | Like 'liftSubExp' but always returns a Regular result with the
-- given expected shape. Reshapes the underlying data if necessary.
liftSubExpRegular ::
  SegLevel ->
  Segments ->
  DistInputs ->
  DistEnv ->
  Shape ->
  SubExp ->
  FlattenM VName
liftSubExpRegular :: SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
expectedShape SubExp
se = do
  case SubExp
se of
    c :: SubExp
c@(Constant PrimValue
_) ->
      Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"lifted_const" (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
$ Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
segmentsShape [SubExp]
segments) SubExp
c)
    Var VName
v -> SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> VName
-> FlattenM VName
liftVarRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
expectedShape VName
v

liftVarRegular ::
  SegLevel ->
  Segments ->
  DistInputs ->
  DistEnv ->
  Shape ->
  VName ->
  FlattenM VName
liftVarRegular :: SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> VName
-> FlattenM VName
liftVarRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
expectedShape VName
x = do
  v <- case VName -> Map VName (Type, ResRep) -> Maybe (Type, ResRep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x (Map VName (Type, ResRep) -> Maybe (Type, ResRep))
-> Map VName (Type, ResRep) -> Maybe (Type, ResRep)
forall a b. (a -> b) -> a -> b
$ DistInputs -> DistEnv -> Map VName (Type, ResRep)
inputReps DistInputs
inps DistEnv
env of
    Just (Type
_, Regular VName
v') -> VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v'
    Just (Type
_, Irregular IrregularRep
irreg) -> do
      rep_dense <- SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl Name
"lifted_irreg" IrregularRep
irreg
      pure $ irregularD rep_dense
    Maybe (Type, ResRep)
Nothing ->
      Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"free_replicated" (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
$ Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
segmentsShape [SubExp]
segments) (VName -> SubExp
Var VName
x)
  v_t <- lookupType v
  if isAcc v_t || arrayShape v_t == expectedShape
    then pure v
    else
      letExp "reg_lifted" . BasicOp $
        Reshape v (reshapeAll (arrayShape v_t) expectedShape)

-- | Can this input array be lifted to a regular array? This holds unless it is
-- represented irregularly. The uniform alternatives lift their inputs regularly
-- (via 'liftSubExpRegular'), which is only valid when the inputs are actually
-- regular.
isRegularInputArr :: DistEnv -> DistInputs -> VName -> Bool
isRegularInputArr :: DistEnv -> DistInputs -> VName -> Bool
isRegularInputArr 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
_) ->
      case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
        Regular {} -> Bool
True
        Irregular {} -> Bool
False
    Maybe DistInput
_ -> Bool
True

liftParam :: (MonadFreshNames m) => SubExp -> FParam SOACS -> m ([FParam GPU], ResRep)
liftParam :: forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> FParam SOACS -> m ([FParam GPU], ResRep)
liftParam SubExp
w FParam SOACS
fparam =
  case Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
FParam SOACS
fparam of
    Prim PrimType
pt -> do
      p <-
        Name -> DeclType -> m (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam
          (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_lifted")
          (TypeBase Shape (ZonkAny 2) -> Shape -> Diet -> DeclType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase Shape (ZonkAny 2)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) Diet
Observe)
      pure ([p], Regular $ paramName p)
    Array PrimType
pt Shape
_ Diet
o -> do
      num_data <-
        Name -> DeclType -> m (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_num_data") (DeclType -> m (Param DeclType)) -> DeclType -> m (Param DeclType)
forall a b. (a -> b) -> a -> b
$ PrimType -> DeclType
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
      segments <-
        newParam (desc <> "_S") $
          arrayOf (Prim int64) (Shape [w]) Observe
      flags <-
        newParam (desc <> "_F") $
          arrayOf (Prim Bool) (Shape [Var (paramName num_data)]) Observe
      offsets <-
        newParam (desc <> "_O") $
          arrayOf (Prim int64) (Shape [w]) Observe
      elems <-
        newParam (desc <> "_D") $
          arrayOf (Prim pt) (Shape [Var (paramName num_data)]) o
      pure
        ( [num_data, segments, flags, offsets, elems],
          Irregular $
            IrregularRep
              { irregularS = paramName segments,
                irregularF = paramName flags,
                irregularO = paramName offsets,
                irregularD = paramName elems,
                irregularK = Dense
              }
        )
    Acc {} ->
      String -> m ([Param DeclType], ResRep)
forall a. HasCallStack => String -> a
error String
"liftParam: Acc"
    Mem {} ->
      String -> m ([Param DeclType], ResRep)
forall a. HasCallStack => String -> a
error String
"liftParam: Mem"
  where
    desc :: Name
desc = VName -> Name
baseName (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
FParam SOACS
fparam)

liftRegularParam :: (MonadFreshNames m) => SubExp -> FParam SOACS -> m (FParam GPU, ResRep)
liftRegularParam :: forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> FParam SOACS -> m (FParam GPU, ResRep)
liftRegularParam SubExp
w FParam SOACS
fparam =
  case Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
FParam SOACS
fparam of
    Prim PrimType
pt -> do
      p <-
        Name -> DeclType -> m (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam
          (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_lifted")
          (TypeBase Shape (ZonkAny 0) -> Shape -> Diet -> DeclType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase Shape (ZonkAny 0)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) Diet
Observe)
      pure (p, Regular $ paramName p)
    Array PrimType
pt Shape
shape Diet
o -> do
      p <-
        Name -> DeclType -> m (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_lifted") (DeclType -> m (Param DeclType)) -> DeclType -> m (Param DeclType)
forall a b. (a -> b) -> a -> b
$
          TypeBase Shape (ZonkAny 1) -> Shape -> Diet -> DeclType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase Shape (ZonkAny 1)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
shape) Diet
o
      pure (p, Regular $ paramName p)
    Acc {} ->
      String -> m (Param DeclType, ResRep)
forall a. HasCallStack => String -> a
error String
"liftParam: Acc"
    Mem {} ->
      String -> m (Param DeclType, ResRep)
forall a. HasCallStack => String -> a
error String
"liftParam: Mem"
  where
    desc :: Name
desc = VName -> Name
baseName (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
FParam SOACS
fparam)

distCerts :: DistInputs -> StmAux a -> DistEnv -> Certs
distCerts :: forall a. DistInputs -> StmAux a -> DistEnv -> Certs
distCerts DistInputs
inps StmAux a
aux DistEnv
env = [VName] -> Certs
Certs ([VName] -> Certs) -> [VName] -> Certs
forall a b. (a -> b) -> a -> b
$ (VName -> VName) -> [VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> VName
f ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Certs -> [VName]
unCerts (Certs -> [VName]) -> Certs -> [VName]
forall a b. (a -> b) -> a -> b
$ StmAux a -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux a
aux
  where
    f :: VName -> VName
f VName
v = case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps of
      Maybe DistInput
Nothing -> VName
v
      Just (DistInputFree VName
vs Type
_) -> VName
vs
      Just (DistInput ResTag
rt Type
_) ->
        case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
          Regular VName
vs -> VName
vs
          Irregular IrregularRep
r -> IrregularRep -> VName
irregularD IrregularRep
r

flattenData :: VName -> FlattenM VName
flattenData :: VName -> FlattenM VName
flattenData VName
vs = do
  t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
vs
  case arrayDims t of
    [SubExp
_] -> VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
vs
    [SubExp]
dims -> do
      n <- Name -> TPrimExp Int64 VName -> FlattenM SubExp
forall (m :: * -> *) a.
(MonadBuilder m, ToExp a) =>
Name -> a -> m SubExp
toSubExp Name
"num_data" (TPrimExp Int64 VName -> FlattenM SubExp)
-> TPrimExp Int64 VName -> FlattenM SubExp
forall a b. (a -> b) -> a -> b
$ [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)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
dims
      letExp (baseName vs <> "_flat") . BasicOp $
        Reshape vs $
          reshapeAll (arrayShape t) (Shape [n])

-- | Only sensible for variables of uniform type.
dataArr :: SegLevel -> Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr :: SegLevel
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
_segments DistEnv
env DistInputs
inps (Var VName
v)
  | Just DistInput
v_inp <- VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps =
      case DistInput
v_inp of
        DistInputFree VName
vs Type
_ -> VName -> FlattenM VName
flattenData VName
vs
        DistInput ResTag
rt Type
_ -> case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
          Irregular IrregularRep
r -> do
            rep_dense <- SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl Name
"dataArr" IrregularRep
r
            pure $ irregularD rep_dense
          Regular VName
vs -> VName -> FlattenM VName
flattenData VName
vs
dataArr SegLevel
_ [SubExp]
segments DistEnv
_ DistInputs
_ SubExp
se = do
  -- The result is a one-dimensional array with one element per segment. With no
  -- enclosing segments there is a single implicit segment, so we replicate over
  -- a unit dimension; replicating over the empty shape would instead yield a
  -- scalar.
  let rep_shape :: Shape
rep_shape = case [SubExp] -> Shape
segmentsShape [SubExp]
segments of
        Shape [] -> [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1]
        Shape
shape -> Shape
shape
  rep <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"rep" (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
$ Shape -> SubExp -> BasicOp
Replicate Shape
rep_shape SubExp
se
  rep_t <- lookupType rep
  let dims = Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
rep_t
  if length dims == 1
    then pure rep
    else do
      n <- toSubExp "n" $ product $ map pe64 dims
      letExp "reshape" $ BasicOp $ Reshape rep $ reshapeAll (arrayShape rep_t) (Shape [n])

-- | Get the irregular representation of a var.
getIrregRep :: SegLevel -> Segments -> DistEnv -> DistInputs -> VName -> FlattenM IrregularRep
getIrregRep :: SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
v =
  case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps of
    Just DistInput
v_inp -> case DistInput
v_inp of
      DistInputFree VName
arr Type
_ -> SegLevel -> [SubExp] -> VName -> FlattenM IrregularRep
mkIrregFromReg SegLevel
lvl [SubExp]
segments VName
arr
      DistInput ResTag
rt Type
_ -> case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
        Irregular IrregularRep
r -> IrregularRep -> FlattenM IrregularRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IrregularRep
r
        Regular VName
arr -> SegLevel -> [SubExp] -> VName -> FlattenM IrregularRep
mkIrregFromReg SegLevel
lvl [SubExp]
segments VName
arr
    Maybe DistInput
Nothing -> do
      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
"_rep") (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
$
          Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
segmentsShape [SubExp]
segments) (VName -> SubExp
Var VName
v)
      mkIrregFromReg lvl segments v'

-- | This function walks through the *unlifted* result types
-- and uses the *lifted* results to construct the corresponding res reps.
--
-- See the 'liftResult' function for the opposite process i.e.
-- turning 'ResRep's into results.
resultToResReps :: [TypeBase s o] -> [VName] -> [ResRep]
resultToResReps :: forall s o. [TypeBase s o] -> [VName] -> [ResRep]
resultToResReps [TypeBase s o]
types [VName]
results =
  ([VName], [ResRep]) -> [ResRep]
forall a b. (a, b) -> b
snd (([VName], [ResRep]) -> [ResRep])
-> ([VName], [ResRep]) -> [ResRep]
forall a b. (a -> b) -> a -> b
$
    ([VName] -> TypeBase s o -> ([VName], ResRep))
-> [VName] -> [TypeBase s o] -> ([VName], [ResRep])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL
      ( \[VName]
rs TypeBase s o
t -> case TypeBase s o
t of
          Prim {} ->
            let (VName
v : [VName]
rs') = [VName]
rs
                rep :: ResRep
rep = VName -> ResRep
Regular VName
v
             in ([VName]
rs', ResRep
rep)
          Array {} ->
            let (VName
_ : VName
segs : VName
flags : VName
offsets : VName
elems : [VName]
rs') = [VName]
rs
                rep :: ResRep
rep = IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep) -> IrregularRep -> ResRep
forall a b. (a -> b) -> a -> b
$ VName -> VName -> VName -> VName -> IrregularKind -> IrregularRep
IrregularRep VName
segs VName
flags VName
offsets VName
elems IrregularKind
Dense
             in ([VName]
rs', ResRep
rep)
          Acc {} -> String -> ([VName], ResRep)
forall a. HasCallStack => String -> a
error String
"resultToResReps: Illegal type 'Acc'"
          Mem {} -> String -> ([VName], ResRep)
forall a. HasCallStack => String -> a
error String
"resultToResReps: Illegal type 'Mem'"
      )
      [VName]
results
      [TypeBase s o]
types

resultToResRepsByDistResult :: [DistResult] -> [VName] -> [ResRep]
resultToResRepsByDistResult :: [DistResult] -> [VName] -> [ResRep]
resultToResRepsByDistResult [DistResult]
dist_res [VName]
results =
  ([VName], [ResRep]) -> [ResRep]
forall a b. (a, b) -> b
snd (([VName], [ResRep]) -> [ResRep])
-> ([VName], [ResRep]) -> [ResRep]
forall a b. (a -> b) -> a -> b
$
    ([VName] -> DistResult -> ([VName], ResRep))
-> [VName] -> [DistResult] -> ([VName], [ResRep])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL
      ( \[VName]
rs DistResult
dist_res' ->
          if DistResult -> Bool
isRegularDistResult DistResult
dist_res'
            then
              let (VName
v : [VName]
rs') = [VName]
rs
               in ([VName]
rs', VName -> ResRep
Regular VName
v)
            else
              let (VName
_ : VName
segs : VName
flags : VName
offsets : VName
elems : [VName]
rs') = [VName]
rs
               in ([VName]
rs', IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep) -> IrregularRep -> ResRep
forall a b. (a -> b) -> a -> b
$ VName -> VName -> VName -> VName -> IrregularKind -> IrregularRep
IrregularRep VName
segs VName
flags VName
offsets VName
elems IrregularKind
Dense)
      )
      [VName]
results
      [DistResult]
dist_res

-- helper to not mess up the tags when generating new ones for the loop parameters
-- probably won't be used in future
localiseInputs :: DistEnv -> DistInputs -> FlattenM (DistInputs, DistEnv, Int)
localiseInputs :: DistEnv -> DistInputs -> FlattenM (DistInputs, DistEnv, Int)
localiseInputs DistEnv
env_outer DistInputs
inps = do
  let step :: (Int, DistEnv)
-> (a, DistInput) -> FlattenM ((Int, DistEnv), (a, DistInput))
step (Int
i, DistEnv
env_acc) (a
v, DistInput
inp) =
        case DistInput
inp of
          DistInputFree 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
i, DistEnv
env_acc), (a
v, VName -> Type -> DistInput
DistInputFree VName
arr Type
t))
          DistInput ResTag
oldrt Type
t -> do
            let newrt :: ResTag
newrt = Int -> ResTag
ResTag Int
i
                rep :: ResRep
rep = ResTag -> DistEnv -> ResRep
resVar ResTag
oldrt DistEnv
env_outer
            env_acc' <- ResTag -> ResRep -> DistEnv -> FlattenM DistEnv
insertRepM ResTag
newrt ResRep
rep DistEnv
env_acc
            pure ((i + 1, env_acc'), (v, DistInput newrt t))

  ((next, env_local), inps_local) <-
    ((Int, DistEnv)
 -> (VName, DistInput)
 -> FlattenM ((Int, DistEnv), (VName, DistInput)))
-> (Int, DistEnv)
-> DistInputs
-> 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, DistInput)
-> FlattenM ((Int, DistEnv), (VName, DistInput))
forall {a}.
(Int, DistEnv)
-> (a, DistInput) -> FlattenM ((Int, DistEnv), (a, DistInput))
step (Int
0, DistEnv
forall a. Monoid a => a
mempty) DistInputs
inps
  pure (inps_local, env_local, next)

-- | Replicate an array to insert new inner dimensions  after the
-- existing segment dimensions.
replicateForDims :: Segments -> Shape -> VName -> FlattenM VName
replicateForDims :: [SubExp] -> Shape -> VName -> FlattenM VName
replicateForDims [SubExp]
segments Shape
dims VName
v = do
  v_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
  let seg_rank = [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
segments
      v_rank = Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
v_t
      dims_rank = Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
dims
      perm = [Int
dims_rank .. Int
dims_rank Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
seg_rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
0 .. Int
dims_rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
seg_rank Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dims_rank .. Int
dims_rank Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
v_rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
  v_rep <-
    letExp (baseName v <> "_reg_rep") . BasicOp $ Replicate dims (Var v)
  letExp (baseName v <> "_reg_rep_tr") . BasicOp $ Rearrange v_rep perm

-- | Flatten a single 'DistStm', producing an updated environment.
flattenDistStm ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistStm ->
  FlattenM DistEnv
flattenDistStm :: FlattenOps -> [SubExp] -> DistEnv -> DistStm -> FlattenM DistEnv
flattenDistStm FlattenOps
ops = FlattenOps
-> FlattenOps -> [SubExp] -> DistEnv -> DistStm -> FlattenM DistEnv
flattenDistStmWith FlattenOps
ops FlattenOps
ops

-- | Flatten a single scalar statement, producing an updated environment.
flattenScalarStm ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Stm SOACS ->
  FlattenM DistEnv
flattenScalarStm :: FlattenOps
-> [SubExp]
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
flattenScalarStm FlattenOps
ops = FlattenOps
-> SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
flattenScalarStmAt FlattenOps
ops (FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops)

-- | 'distributeBody' with the settings in the given 'FlattenOps'.
distributeBodyWith ::
  FlattenOps ->
  Scope rep ->
  Segments ->
  DistInputs ->
  Body SOACS ->
  (DistInputs, DistStms)
distributeBodyWith :: forall rep.
FlattenOps
-> Scope rep
-> [SubExp]
-> DistInputs
-> Body SOACS
-> (DistInputs, DistStms)
distributeBodyWith FlattenOps
ops =
  DistIrregularity
-> FunHasParallelism
-> Scope rep
-> [SubExp]
-> DistInputs
-> Body SOACS
-> (DistInputs, DistStms)
forall rep.
DistIrregularity
-> FunHasParallelism
-> Scope rep
-> [SubExp]
-> DistInputs
-> Body SOACS
-> (DistInputs, DistStms)
distributeBody (FlattenOps -> DistIrregularity
flattenIrregularity FlattenOps
ops) (FlattenOps -> FunHasParallelism
flattenFunHasParallelism FlattenOps
ops)

-- | 'distributeMap' with the settings in the given 'FlattenOps'.
distributeMapWith ::
  FlattenOps ->
  Scope rep ->
  Pat Type ->
  Segments ->
  [MapArray t] ->
  Lambda SOACS ->
  (Distributed, M.Map ResTag t)
distributeMapWith :: forall rep t.
FlattenOps
-> Scope rep
-> Pat Type
-> [SubExp]
-> [MapArray t]
-> Lambda SOACS
-> (Distributed, Map ResTag t)
distributeMapWith FlattenOps
ops =
  DistIrregularity
-> FunHasParallelism
-> Scope rep
-> Pat Type
-> [SubExp]
-> [MapArray t]
-> Lambda SOACS
-> (Distributed, Map ResTag t)
forall rep t.
DistIrregularity
-> FunHasParallelism
-> Scope rep
-> Pat Type
-> [SubExp]
-> [MapArray t]
-> Lambda SOACS
-> (Distributed, Map ResTag t)
distributeMap (FlattenOps -> DistIrregularity
flattenIrregularity FlattenOps
ops) (FlattenOps -> FunHasParallelism
flattenFunHasParallelism FlattenOps
ops)

-- | Continue flattening at the given seg level, adjusting the irregularity
-- handling mode to match. Intrablock code cannot use the machinery for
-- flattening irregular arrays, as it produces SegOps whose sizes are bound
-- inside the enclosing kernel. 'SequentialiseIrregularAll' is requested by the
-- user rather than implied by the level, so we keep it if provided.
atSegLevel :: SegLevel -> FlattenOps -> FlattenOps
atSegLevel :: SegLevel -> FlattenOps -> FlattenOps
atSegLevel SegLevel
lvl FlattenOps
ops =
  FlattenOps
ops {flattenSegLevel = lvl, flattenIrregularity = irreg}
  where
    irreg :: DistIrregularity
irreg = case (FlattenOps -> DistIrregularity
flattenIrregularity FlattenOps
ops, SegLevel
lvl) of
      (DistIrregularity
SequentialiseIrregularAll, SegLevel
_) -> DistIrregularity
SequentialiseIrregularAll
      (DistIrregularity
_, SegThreadInBlock {}) -> DistIrregularity
SequentialiseIrregularBasicOps
      (DistIrregularity, SegLevel)
_ -> DistIrregularity
DistributeIrregular