module Futhark.Pass.Flatten.BasicOp (flattenBasicOp) where

import Control.Monad
import Data.Foldable
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Tuple.Solo
import Futhark.IR.GPU
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Util.IntegralExp
import Prelude hiding (div, quot, rem)

-- Do 'map2 (++) A B' where 'A' and 'B' are irregular arrays and have the same
-- number of subarrays
concatIrreg ::
  SegLevel ->
  Segments ->
  DistEnv ->
  VName ->
  [IrregularRep] ->
  FlattenM IrregularRep
concatIrreg :: SegLevel
-> [SubExp]
-> DistEnv
-> VName
-> [IrregularRep]
-> FlattenM IrregularRep
concatIrreg SegLevel
lvl [SubExp]
_segments DistEnv
_env VName
ns [IrregularRep]
reparr = do
  -- Concatenation does not change the number of segments - it simply
  -- makes each of them larger.

  num_segments <- Int -> Type -> SubExp
forall o. Int -> TypeBase Shape 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
ns

  -- Constructs the full list size / shape that should hold the final results.
  ns_full <- letExp (baseName ns <> "_full") <=< segMap lvl (MkSolo num_segments) $
    \(MkSolo SubExp
i) -> do
      old_segments <-
        [IrregularRep]
-> (IrregularRep -> FlattenM SubExp) -> FlattenM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [IrregularRep]
reparr ((IrregularRep -> FlattenM SubExp) -> FlattenM [SubExp])
-> (IrregularRep -> FlattenM SubExp) -> FlattenM [SubExp]
forall a b. (a -> b) -> a -> b
$ \IrregularRep
rep ->
          Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"old_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 (IrregularRep -> VName
irregularS IrregularRep
rep) [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
      new_segment <-
        letSubExp "new_segment"
          =<< toExp (sum $ map pe64 old_segments)
      pure $ subExpsRes [new_segment]

  (ns_full_F, ns_full_O, _ns_II1) <- doRepIota lvl ns_full

  repIota <- mapM (doRepIota lvl . irregularS) reparr
  segIota <- mapM (doSegIota lvl . irregularS) reparr

  let (_, _, rep_II1) = unzip3 repIota
  let (_, _, rep_II2) = unzip3 segIota

  n_arr <- mapM (fmap (arraySize 0) . lookupType) rep_II1

  -- Calculate offsets for the scatter operations
  let shapes = (IrregularRep -> VName) -> [IrregularRep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map IrregularRep -> VName
irregularS [IrregularRep]
reparr
  scatter_offsets <-
    letTupExp "irregular_scatter_offsets" <=< segMap lvl (MkSolo num_segments) $
      \(MkSolo SubExp
i) -> do
        segment_sizes <-
          [VName] -> (VName -> FlattenM SubExp) -> FlattenM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
shapes ((VName -> FlattenM SubExp) -> FlattenM [SubExp])
-> (VName -> FlattenM SubExp) -> FlattenM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
shape ->
            Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment_size" (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 (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
        let scanned = (TPrimExp Int64 VName
 -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> TPrimExp Int64 VName
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
(+) TPrimExp Int64 VName
0 ([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]
segment_sizes
        sumprefix <- mapM (letSubExp "segment_prefix" <=< toExp) (init scanned)
        pure $ subExpsRes sumprefix

  scatter_offsets_T <-
    letTupExp "irregular_scatter_offsets_T" <=< segMap lvl (MkSolo num_segments) $
      \(MkSolo SubExp
i) -> do
        columns <-
          [VName] -> (VName -> FlattenM SubExp) -> FlattenM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
scatter_offsets ((VName -> FlattenM SubExp) -> FlattenM [SubExp])
-> (VName -> FlattenM SubExp) -> FlattenM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
offsets ->
            Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment_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
offsets [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
        pure $ subExpsRes columns

  m <- arraySize 0 <$> lookupType ns_full_F
  data_t <- lookupType (irregularD (head reparr))
  let pt = Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
data_t
  let result_type = PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
m]) NoMode
NoMode
  elems_blank <- letExp "blank_res" =<< eBlank result_type

  -- Scatter data into result array
  elems <-
    foldlM
      ( \VName
elems (IrregularRep
reparr1, VName
scatter_offset, SubExp
n, VName
ii1, VName
ii2) -> do
          Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"irregular_scatter_elems" (Exp GPU -> FlattenM VName)
-> ((SubExp -> FlattenM (VName, SubExp)) -> FlattenM (Exp GPU))
-> (SubExp -> FlattenM (VName, SubExp))
-> FlattenM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> VName
-> SubExp
-> (SubExp -> FlattenM (VName, SubExp))
-> FlattenM (Exp GPU)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName -> SubExp -> (SubExp -> m (VName, SubExp)) -> m (Exp GPU)
genScatter SegLevel
lvl VName
elems SubExp
n ((SubExp -> FlattenM (VName, SubExp)) -> FlattenM VName)
-> (SubExp -> FlattenM (VName, SubExp)) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ \SubExp
gid -> do
            -- Which segment we are in.
            segment_i <-
              Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment_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
ii1 [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gid]

            -- Get segment offset in final array
            segment_o <-
              letSubExp "segment_o" =<< eIndex ns_full_O [eSubExp segment_i]

            -- Get local segment offset
            segment_local_o <-
              letSubExp "segment_local_o"
                =<< eIndex scatter_offset [eSubExp segment_i]

            o' <- letSubExp "o" =<< eIndex ii2 [eSubExp gid]
            src_segment_o <-
              letSubExp "src_segment_o" =<< eIndex (irregularO reparr1) [eSubExp segment_i]
            src_i <-
              letSubExp "src_i" <=< toExp $ pe64 src_segment_o + pe64 o'
            v' <-
              letSubExp "v" =<< eIndex (irregularD reparr1) [eSubExp src_i]

            -- Index to write `v'` at
            i <-
              letExp "i" =<< toExp (pe64 o' + pe64 segment_local_o + pe64 segment_o)

            pure (i, v')
      )
      elems_blank
      $ L.zip5 reparr scatter_offsets_T n_arr rep_II1 rep_II2

  pure $
    IrregularRep
      { irregularS = ns_full,
        irregularF = ns_full_F,
        irregularO = ns_full_O,
        irregularD = elems,
        irregularK = Dense
      }

-- We also can do reearange -> concat -> rearrange but this should be more efficient
concatIrregAlongDim ::
  SegLevel ->
  Segments ->
  DistEnv ->
  VName ->
  [IrregularRep] ->
  [Type] ->
  DistInputs ->
  Int ->
  FlattenM IrregularRep
concatIrregAlongDim :: SegLevel
-> [SubExp]
-> DistEnv
-> VName
-> [IrregularRep]
-> [Type]
-> DistInputs
-> Int
-> FlattenM IrregularRep
concatIrregAlongDim SegLevel
lvl [SubExp]
segments DistEnv
env VName
ns [IrregularRep]
rep_arr [Type]
type_arr DistInputs
inps Int
d = do
  num_segments <- Int -> Type -> SubExp
forall o. Int -> TypeBase Shape 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
ns

  ns_full <- letExp (baseName ns <> "_full") <=< segMap lvl (MkSolo num_segments) $
    \(MkSolo SubExp
i) -> do
      old_segments <-
        [IrregularRep]
-> (IrregularRep -> FlattenM SubExp) -> FlattenM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [IrregularRep]
rep_arr ((IrregularRep -> FlattenM SubExp) -> FlattenM [SubExp])
-> (IrregularRep -> FlattenM SubExp) -> FlattenM [SubExp]
forall a b. (a -> b) -> a -> b
$ \IrregularRep
rep ->
          Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"old_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 (IrregularRep -> VName
irregularS IrregularRep
rep) [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
      new_segment <-
        letSubExp "new_segment"
          =<< toExp (sum $ map pe64 old_segments)
      pure $ subExpsRes [new_segment]

  (ns_full_F, ns_full_O, _ns_II1) <- doRepIota lvl ns_full

  repIota <- mapM (doRepIota lvl . irregularS) rep_arr
  segIota <- mapM (doSegIota lvl . irregularS) rep_arr

  let (_, _, rep_II1) = unzip3 repIota
  let (_, _, rep_II2) = unzip3 segIota

  n_arr <- mapM (fmap (arraySize 0) . lookupType) rep_II1

  scatter_info <-
    letTupExp "irregular_scatter_offsets" <=< segMap lvl (MkSolo num_segments) $
      \(MkSolo SubExp
i) -> do
        seg_is <- [SubExp] -> SubExp -> FlattenM [SubExp]
segmentCoordsFromFlat [SubExp]
segments SubExp
i

        block_sizes <-
          forM type_arr $ \Type
t -> do
            v_dims <- [SubExp]
-> DistEnv -> [SubExp] -> DistInputs -> Type -> FlattenM [SubExp]
forall o.
[SubExp]
-> DistEnv
-> [SubExp]
-> DistInputs
-> TypeBase Shape o
-> FlattenM [SubExp]
readTypeDims [SubExp]
segments DistEnv
env [SubExp]
seg_is DistInputs
inps Type
t
            letSubExp "block_size" =<< toExp (product $ map pe64 $ drop d v_dims)

        let scanned = (TPrimExp Int64 VName
 -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> TPrimExp Int64 VName
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
(+) TPrimExp Int64 VName
0 ([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]
block_sizes
        sum_prefix <- mapM (letSubExp "segment_prefix" <=< toExp) (init scanned)
        total_block <- letSubExp "total_block" =<< toExp (last scanned)

        pure $ subExpsRes (block_sizes <> sum_prefix <> [total_block])

  let k = [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
type_arr
      (scatter_blocks, rest) = splitAt k scatter_info
      (scatter_offsets, [total_block_size]) = splitAt k rest

  m <- arraySize 0 <$> lookupType ns_full_F
  data_t <- lookupType (irregularD (head rep_arr))
  let pt = Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
data_t
  let result_type = PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
m]) NoMode
NoMode
  elems_blank <- letExp "blank_res" =<< eBlank result_type

  -- Scatter data into result array
  elems <-
    foldlM
      ( \VName
elems (IrregularRep
reparr1, VName
scatter_block, VName
scatter_offset, SubExp
n, VName
ii1, VName
ii2) -> do
          Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"irregular_scatter_elems" (Exp GPU -> FlattenM VName)
-> ((SubExp -> FlattenM (VName, SubExp)) -> FlattenM (Exp GPU))
-> (SubExp -> FlattenM (VName, SubExp))
-> FlattenM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> VName
-> SubExp
-> (SubExp -> FlattenM (VName, SubExp))
-> FlattenM (Exp GPU)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName -> SubExp -> (SubExp -> m (VName, SubExp)) -> m (Exp GPU)
genScatter SegLevel
lvl VName
elems SubExp
n ((SubExp -> FlattenM (VName, SubExp)) -> FlattenM VName)
-> (SubExp -> FlattenM (VName, SubExp)) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ \SubExp
gid -> do
            -- Which segment we are in.
            segment_i <-
              Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment_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
ii1 [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gid]

            -- Get segment offset in final array
            segment_o <-
              letSubExp "segment_o" =<< eIndex ns_full_O [eSubExp segment_i]

            -- Get local segment offset
            segment_local_o <-
              letSubExp "segment_local_o"
                =<< eIndex scatter_offset [eSubExp segment_i]

            o' <- letSubExp "o" =<< eIndex ii2 [eSubExp gid]
            src_segment_o <-
              letSubExp "src_segment_o" =<< eIndex (irregularO reparr1) [eSubExp segment_i]
            src_i <-
              letSubExp "src_i" <=< toExp $ pe64 src_segment_o + pe64 o'
            v' <-
              letSubExp "v" =<< eIndex (irregularD reparr1) [eSubExp src_i]

            scatter_block_size <-
              letSubExp "scatter_block_size" =<< eIndex scatter_block [eSubExp segment_i]

            scatter_total_block_size <-
              letSubExp "scatter_total_block_size" =<< eIndex total_block_size [eSubExp segment_i]

            outer_i <-
              letSubExp "outer_i" =<< toExp (pe64 o' `div` pe64 scatter_block_size)

            i <-
              letExp "i"
                =<< toExp
                  ( pe64 o'
                      + pe64 outer_i * (pe64 scatter_total_block_size - pe64 scatter_block_size)
                      + pe64 segment_local_o
                      + pe64 segment_o
                  )
            pure (i, v')
      )
      elems_blank
      $ L.zip6 rep_arr scatter_blocks scatter_offsets n_arr rep_II1 rep_II2

  pure $
    IrregularRep
      { irregularS = ns_full,
        irregularF = ns_full_F,
        irregularO = ns_full_O,
        irregularD = elems,
        irregularK = Dense
      }

-- Do 'map2 replicate ns A', where 'A' is an irregular array (and so
-- is the result, obviously).
replicateIrreg ::
  SegLevel ->
  Segments ->
  DistEnv ->
  VName ->
  Name ->
  IrregularRep ->
  FlattenM IrregularRep
replicateIrreg :: SegLevel
-> [SubExp]
-> DistEnv
-> VName
-> Name
-> IrregularRep
-> FlattenM IrregularRep
replicateIrreg SegLevel
lvl [SubExp]
_segments DistEnv
_env VName
ns Name
desc IrregularRep
rep = do
  -- Replication does not change the number of segments - it simply
  -- makes each of them larger.

  num_segments <- Int -> Type -> SubExp
forall o. Int -> TypeBase Shape 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
ns

  -- ns multipled with existing segment sizes.
  ns_full <- letExp (baseName ns <> "_full") <=< segMap lvl (MkSolo num_segments) $
    \(MkSolo SubExp
i) -> do
      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
=<< VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
ns [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
      old_segment <-
        letSubExp "old_segment" =<< eIndex (irregularS rep) [eSubExp i]
      full_segment <-
        letSubExp "new_segment" =<< toExp (pe64 n * pe64 old_segment)
      pure $ subExpsRes [full_segment]

  (ns_full_F, ns_full_O, ns_full_D) <- doRepIota lvl ns_full
  (_, _, flat_to_segs) <- doSegIota lvl ns_full

  w <- arraySize 0 <$> lookupType ns_full_D

  elems <- letExp (desc <> "_rep_D") <=< segMap lvl (MkSolo w) $ \(MkSolo SubExp
i) -> do
    -- Which segment we are in.
    segment_i <-
      Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment_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
ns_full_D [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    -- Size of original segment.
    old_segment <-
      letSubExp "old_segment" =<< eIndex (irregularS rep) [eSubExp segment_i]
    -- Index of value inside *new* segment.
    j_new <-
      letSubExp "j_new" =<< eIndex flat_to_segs [eSubExp i]
    -- Index of value inside *old* segment.
    j_old <-
      letSubExp "j_old" =<< toExp (pe64 j_new `rem` pe64 old_segment)
    -- Offset of values in original segment.
    offset <-
      letSubExp "offset" =<< eIndex (irregularO rep) [eSubExp segment_i]
    v <-
      letSubExp "v"
        =<< eIndex (irregularD rep) [toExp $ pe64 offset + pe64 j_old]
    pure $ subExpsRes [v]

  pure $
    IrregularRep
      { irregularS = ns_full,
        irregularF = ns_full_F,
        irregularO = ns_full_O,
        irregularD = elems,
        irregularK = Dense
      }

rearrangeFlat :: (IntegralExp num) => [Int] -> [num] -> num -> num
rearrangeFlat :: forall num. IntegralExp num => [Int] -> [num] -> num -> num
rearrangeFlat [Int]
perm [num]
dims num
i =
  [num] -> [num] -> num
forall num. IntegralExp num => [num] -> [num] -> num
flattenIndex [num]
dims ([num] -> num) -> [num] -> num
forall a b. (a -> b) -> a -> b
$
    [Int] -> [num] -> [num]
forall a. [Int] -> [a] -> [a]
rearrangeShape ([Int] -> [Int]
rearrangeInverse [Int]
perm) ([num] -> [num]) -> [num] -> [num]
forall a b. (a -> b) -> a -> b
$
      [num] -> num -> [num]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex ([Int] -> [num] -> [num]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm [num]
dims) num
i

segmentCoordsFromFlat :: Segments -> SubExp -> FlattenM [SubExp]
segmentCoordsFromFlat :: [SubExp] -> SubExp -> FlattenM [SubExp]
segmentCoordsFromFlat [SubExp]
segments SubExp
seg_i =
  (TPrimExp Int64 VName -> FlattenM SubExp)
-> [TPrimExp Int64 VName] -> 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 (Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"seg_coord" (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
$
    [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex ((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
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (Shape -> [SubExp]) -> Shape -> [SubExp]
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Shape
segmentsShape [SubExp]
segments) (SubExp -> TPrimExp Int64 VName
pe64 SubExp
seg_i)

-- TODO: We do not need to actually make this Dense
rearrangeIrreg ::
  SegLevel ->
  Segments ->
  DistEnv ->
  DistInputs ->
  TypeBase Shape o ->
  [Int] ->
  IrregularRep ->
  FlattenM IrregularRep
rearrangeIrreg :: forall o.
SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> TypeBase Shape o
-> [Int]
-> IrregularRep
-> FlattenM IrregularRep
rearrangeIrreg SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps TypeBase Shape o
v_t [Int]
perm IrregularRep
ir = do
  (IrregularRep shape _ offsets elems _) <- SegLevel -> IrregularRep -> FlattenM IrregularRep
flattenIrregularRep SegLevel
lvl IrregularRep
ir
  (new_F, new_O, ii1_vss) <- doRepIota lvl shape
  (_, _, ii2_vss) <- doSegIota lvl shape
  m <- arraySize 0 <$> lookupType ii1_vss
  elems' <- letExp "elems_rearrange" <=< renameExp <=< segMap lvl (MkSolo m) $
    \(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
ii1_vss [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
      offset <- letSubExp "offset" =<< eIndex offsets [eSubExp seg_i]
      in_seg_i <- letSubExp "in_seg_i" =<< eIndex ii2_vss [eSubExp i]
      seg_is <- segmentCoordsFromFlat segments seg_i
      v_dims <- readTypeDims segments env seg_is inps v_t
      let v_dims' = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
v_dims
          in_seg_is_tr = [Int]
-> [TPrimExp Int64 VName]
-> TPrimExp Int64 VName
-> TPrimExp Int64 VName
forall num. IntegralExp num => [Int] -> [num] -> num -> num
rearrangeFlat [Int]
perm [TPrimExp Int64 VName]
v_dims' (TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
in_seg_i
      v' <-
        letSubExp "v"
          =<< eIndex elems [toExp $ pe64 offset + in_seg_is_tr]
      pure [subExpRes v']
  pure $
    IrregularRep
      { irregularS = shape,
        irregularF = new_F,
        irregularO = new_O,
        irregularD = elems',
        irregularK = Dense
      }

-- | Input common to 'flattenBasicOp': the segment level, the enclosing
-- segments, the distribution environment and inputs, the result being produced,
-- and the statement auxiliary information.
data TrCtx = TrCtx SegLevel Segments DistEnv DistInputs DistResult (StmAux ())

transformArrayLit ::
  TrCtx ->
  [SubExp] ->
  Type ->
  FlattenM DistEnv
-- Potentially no need for this case.
transformArrayLit :: TrCtx -> [SubExp] -> Type -> FlattenM DistEnv
transformArrayLit (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) [] Type
row_type
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
row_type) = do
      let resultType :: Type
resultType =
            PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array
              (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
row_type)
              ([SubExp] -> Shape
segmentsShape [SubExp]
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
row_type)
              NoMode
NoMode
      v <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"arraylit_empty_reg" (Exp GPU -> FlattenM VName) -> FlattenM (Exp GPU) -> FlattenM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank Type
resultType
      pure $ insertRegulars [distResTag res] [v] env
  | Bool
otherwise = do
      ns <- SegLevel
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps (SubExp -> FlattenM VName) -> SubExp -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
      (flags, offsets, _elems) <- doRepIota lvl ns
      let resultType = PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
row_type) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0]) NoMode
NoMode
      elems <- letExp "arraylit_empty_elems" =<< eBlank resultType
      insertIrregularM ns flags offsets (distResTag res) elems Dense env
transformArrayLit (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) [SubExp]
vs Type
row_type
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
row_type) = do
      res_v <-
        if (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) [SubExp]
vs
          then do
            let seg_shape :: Shape
seg_shape = [SubExp] -> Shape
segmentsShape [SubExp]
segments
                one :: SubExp
one = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1
                arr_outer_dim :: SubExp
arr_outer_dim = IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
vs
                expected :: Shape
expected = Shape
seg_shape Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
row_type
                stacked :: Shape
stacked = Shape
seg_shape Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
one] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
row_type
                d :: Int
d = [SubExp] -> Int
segmentsRank [SubExp]
segments

            vs_reg <- (SubExp -> FlattenM VName) -> [SubExp] -> FlattenM [VName]
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
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
expected) [SubExp]
vs

            vs_reg_1 <-
              forM vs_reg $ \VName
v -> do
                v_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
                letExp (baseName v <> "_stack") $
                  BasicOp $
                    Reshape v $
                      reshapeAll (arrayShape v_t) stacked

            case vs_reg_1 of
              [] -> String -> FlattenM VName
forall a. HasCallStack => String -> a
error String
"flattenBasicOp: empty ArrayLit cannot have variant elements"
              [VName
v] ->
                VName -> FlattenM VName
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
              VName
v : [VName]
vs' ->
                Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"arraylit_reg" (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
$ Int -> NonEmpty VName -> SubExp -> BasicOp
Concat Int
d (VName
v VName -> [VName] -> NonEmpty VName
forall a. a -> [a] -> NonEmpty a
NE.:| [VName]
vs') SubExp
arr_outer_dim
          else do
            base_v <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"arraylit_base" (Exp (Rep FlattenM) -> FlattenM VName)
-> Exp (Rep FlattenM) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Type -> BasicOp
ArrayLit [SubExp]
vs Type
row_type
            letExp "arraylit_reg" $
              BasicOp $
                Replicate (segmentsShape segments) (Var base_v)
      pure $ insertRegulars [distResTag res] [res_v] env
  | Bool
otherwise = do
      let arr_outer_dim :: SubExp
arr_outer_dim = IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
vs
      vs_reparr <- (SubExp -> FlattenM VName) -> [SubExp] -> FlattenM [VName]
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
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps) [SubExp]
vs
      dim_arrs <- mapM (dataArr lvl segments env inps) (arrayDims row_type)
      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)
      ~[row_size, full_size] <- letTupExp "arraylit_row_size" <=< segMap lvl (MkSolo num_segments) $ \(MkSolo SubExp
i) -> do
        vals <- (VName -> FlattenM SubExp) -> [VName] -> 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 (\VName
dim_arr -> Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"dim_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
dim_arr [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]) [VName]
dim_arrs
        n <- letSubExp "n" <=< toExp $ product $ map pe64 vals
        fs <- letSubExp "fs" <=< toExp $ pe64 n * pe64 arr_outer_dim
        pure $ subExpsRes [n, fs]

      (_, _, row_II1) <- doRepIota lvl row_size
      (_, _, row_II2) <- doSegIota lvl row_size

      row_flat_size <- arraySize 0 <$> lookupType row_II1

      (full_flags, full_offset, full_II1) <- doRepIota lvl full_size

      m <- arraySize 0 <$> lookupType full_II1
      let pt = Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
row_type
      let resultType = PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
m]) NoMode
NoMode
      elems_blank <- letExp "blank_res" =<< eBlank resultType

      elems <-
        foldlM
          ( \VName
elems (Integer
var_num, VName
arr) -> do
              Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"irregular_scatter_elems" (Exp GPU -> FlattenM VName)
-> ((SubExp -> FlattenM (VName, SubExp)) -> FlattenM (Exp GPU))
-> (SubExp -> FlattenM (VName, SubExp))
-> FlattenM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> VName
-> SubExp
-> (SubExp -> FlattenM (VName, SubExp))
-> FlattenM (Exp GPU)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName -> SubExp -> (SubExp -> m (VName, SubExp)) -> m (Exp GPU)
genScatter SegLevel
lvl VName
elems SubExp
row_flat_size ((SubExp -> FlattenM (VName, SubExp)) -> FlattenM VName)
-> (SubExp -> FlattenM (VName, SubExp)) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ \SubExp
gid -> do
                -- Which segment we are in.
                segment_i <-
                  Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segment_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
row_II1 [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gid]

                row_size_i <-
                  letSubExp "row_size_i" =<< eIndex row_size [eSubExp segment_i]

                segment_global_o <-
                  letSubExp "segment_global_o"
                    =<< eIndex full_offset [eSubExp segment_i]

                v' <-
                  letSubExp "v" =<< eIndex arr [eSubExp gid]

                o' <- letSubExp "o" =<< eIndex row_II2 [eSubExp gid]

                i <-
                  letExp "i"
                    =<< toExp
                      ( pe64 o'
                          + pe64 segment_global_o
                          + pe64 row_size_i * pe64 (intConst Int64 var_num)
                      )

                pure (i, v')
          )
          elems_blank
          $ zip [0 ..] vs_reparr

      insertIrregularM full_size full_flags full_offset (distResTag res) elems Dense env

transformArrayVal ::
  TrCtx ->
  [PrimValue] ->
  PrimType ->
  FlattenM DistEnv
transformArrayVal :: TrCtx -> [PrimValue] -> PrimType -> FlattenM DistEnv
transformArrayVal (TrCtx SegLevel
_lvl [SubExp]
segments DistEnv
env DistInputs
_inps DistResult
res StmAux ()
_aux) [PrimValue]
vs PrimType
row_type = do
  base_v <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"arraylit_base" (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
$ [PrimValue] -> PrimType -> BasicOp
ArrayVal [PrimValue]
vs PrimType
row_type
  res_v <- letExp "arraylit_reg" $ BasicOp $ Replicate (segmentsShape segments) (Var base_v)
  pure $ insertRegulars [distResTag res] [res_v] env

transformReshape ::
  TrCtx ->
  VName ->
  NewShape SubExp ->
  FlattenM DistEnv
transformReshape :: TrCtx -> VName -> NewShape SubExp -> FlattenM DistEnv
transformReshape (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) VName
arr NewShape SubExp
reshape
  | DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> NewShape SubExp -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) NewShape SubExp
reshape) = do
      let outer :: Shape
outer = [SubExp] -> Shape
segmentsShape [SubExp]
segments
          inner_target :: Shape
inner_target = NewShape SubExp -> Shape
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
reshape
          reshape' :: NewShape SubExp
reshape' = Shape -> NewShape SubExp
forall new. ShapeBase new -> NewShape new
reshapeCoerce Shape
outer NewShape SubExp -> NewShape SubExp -> NewShape SubExp
forall a. Semigroup a => a -> a -> a
<> Shape -> NewShape SubExp -> NewShape SubExp
newshapeInner Shape
outer NewShape SubExp
reshape

      arr_t <- DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
arr
      let arr_shape = Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
arr_t
      let unform_arr = Bool -> Bool
not ((SubExp -> Bool) -> Shape -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) Shape
arr_shape)
      if unform_arr
        then do
          arr' <-
            liftSubExpRegular
              lvl
              segments
              inps
              env
              (outer <> arr_shape)
              (Var arr)
          v <- certifying (distCerts inps aux env) . letExp "reshape_reg" . BasicOp $ Reshape arr' reshape'
          pure $ insertRegulars [distResTag res] [v] env
        else do
          arr' <-
            liftSubExpRegular
              lvl
              segments
              inps
              env
              (outer <> inner_target)
              (Var arr)
          pure $ insertRegulars [distResTag res] [arr'] env
  | Bool
otherwise = do
      irreg_v <- SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
arr
      insertRepM (distResTag res) (Irregular irreg_v) env

transformIndex ::
  TrCtx ->
  VName ->
  Slice SubExp ->
  FlattenM DistEnv
transformIndex :: TrCtx -> VName -> Slice SubExp -> FlattenM DistEnv
transformIndex (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) VName
arr Slice SubExp
slice
  | -- Must be a regular result...
    DistResult -> Bool
isRegularDistResult DistResult
res,
    -- And a regular input...
    DistEnv -> DistInputs -> VName -> Bool
isRegularInputArr DistEnv
env DistInputs
inps VName
arr,
    -- And uniform slices.
    Bool -> Bool
not ((SubExp -> Bool) -> Slice SubExp -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) Slice SubExp
slice) = do
      arr_t <- DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
arr
      arr' <-
        liftSubExpRegular
          lvl
          segments
          inps
          env
          (segmentsShape segments <> arrayShape arr_t)
          (Var arr)
      let segmentSlice = (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim ([SubExp] -> [DimIndex SubExp])
-> ([SubExp] -> [SubExp]) -> [SubExp] -> [DimIndex SubExp]
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
      v <-
        certifying (distCerts inps aux env) . letExp "index_reg" . BasicOp $
          Index arr' (Slice $ segmentSlice segments <> unSlice slice)
      pure $ insertRegulars [distResTag res] [v] env
  | DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) (Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice)) = do
      let space :: [SubExp]
space = Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ([SubExp] -> Shape
segmentsShape [SubExp]
segments) [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
      v <-
        Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"index_reg_gather"
          (Exp GPU -> FlattenM VName)
-> ((NonEmpty SubExp -> FlattenM Result) -> FlattenM (Exp GPU))
-> (NonEmpty SubExp -> FlattenM Result)
-> FlattenM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Exp GPU -> FlattenM (Exp GPU)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Exp rep -> m (Exp rep)
renameExp
          (Exp GPU -> FlattenM (Exp GPU))
-> ((NonEmpty SubExp -> FlattenM Result) -> FlattenM (Exp GPU))
-> (NonEmpty SubExp -> FlattenM Result)
-> FlattenM (Exp GPU)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> NonEmpty SubExp
-> (NonEmpty 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] -> NonEmpty SubExp
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList [SubExp]
space)
          ((NonEmpty SubExp -> FlattenM Result) -> FlattenM VName)
-> (NonEmpty SubExp -> FlattenM Result) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ \NonEmpty SubExp
is -> do
            let ([SubExp]
seg_is, [SubExp]
in_is) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
segmentsRank [SubExp]
segments) (NonEmpty SubExp -> [SubExp]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty SubExp
is)
            [SubExp] -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()
readInputs [SubExp]
segments DistEnv
env [SubExp]
seg_is DistInputs
inps
            let slice' :: [TPrimExp Int64 VName]
slice' = Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice ((SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice) ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
in_is)
            StmAux () -> FlattenM Result -> FlattenM Result
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
aux (FlattenM Result -> FlattenM Result)
-> FlattenM Result -> FlattenM Result
forall a b. (a -> b) -> a -> b
$
              (SubExp -> Result) -> FlattenM SubExp -> FlattenM Result
forall a b. (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([SubExp] -> Result
subExpsRes ([SubExp] -> Result) -> (SubExp -> [SubExp]) -> SubExp -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure) (FlattenM SubExp -> FlattenM Result)
-> (Exp GPU -> FlattenM SubExp) -> Exp GPU -> FlattenM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"v"
                (Exp GPU -> FlattenM Result)
-> FlattenM (Exp GPU) -> FlattenM Result
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 GPU))
-> [TPrimExp Int64 VName] -> [FlattenM (Exp GPU)]
forall a b. (a -> b) -> [a] -> [b]
map 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]
slice')
      pure $ insertRegulars [distResTag res] [v] env
  | Bool
otherwise = do
      -- Maximally nonuniform case.
      num_segments <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"num_segments" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp ([SubExp] -> TPrimExp Int64 VName
segmentCount [SubExp]
segments)
      ns <- letExp "slice_sizes" <=< renameExp <=< segMap lvl (MkSolo num_segments) $ \(MkSolo SubExp
segment) -> do
        segment_is <- [SubExp] -> SubExp -> FlattenM [SubExp]
segmentCoordsFromFlat [SubExp]
segments SubExp
segment
        slice_ns <- mapM (readInput segments env segment_is inps) $ sliceDims slice
        fmap varsRes . letTupExp "n" <=< toExp $ product $ map pe64 slice_ns
      (_n, offsets, m) <- exScanAndSum lvl ns
      (_, _, repiota_D) <- doRepIota lvl ns
      flags <- genFlags lvl m offsets
      elems <- letExp "index_irreg_elems" <=< renameExp <=< segMap lvl (NE.singleton m) $ \NonEmpty SubExp
is -> 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
repiota_D (NonEmpty (FlattenM (Exp (Rep FlattenM)))
-> [FlattenM (Exp (Rep FlattenM))]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (NonEmpty (FlattenM (Exp (Rep FlattenM)))
 -> [FlattenM (Exp (Rep FlattenM))])
-> NonEmpty (FlattenM (Exp (Rep FlattenM)))
-> [FlattenM (Exp (Rep FlattenM))]
forall a b. (a -> b) -> a -> b
$ (SubExp -> FlattenM (Exp (Rep FlattenM)))
-> NonEmpty SubExp -> NonEmpty (FlattenM (Exp (Rep FlattenM)))
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp NonEmpty SubExp
is)
        segment_start <- letSubExp "segment_start" =<< eIndex offsets [eSubExp segment]
        segment_is <- segmentCoordsFromFlat segments segment
        readInputs segments env segment_is inps
        let slice' =
              Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice ((SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice) ([TPrimExp Int64 VName] -> [TPrimExp Int64 VName])
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$
                [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 (Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice)) (TPrimExp Int64 VName -> [TPrimExp Int64 VName])
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$
                  TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
subtract (SubExp -> TPrimExp Int64 VName
pe64 SubExp
segment_start) (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
pe64 (SubExp -> TPrimExp Int64 VName) -> SubExp -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$
                    NonEmpty SubExp -> SubExp
forall a. NonEmpty a -> a
NE.head NonEmpty SubExp
is
        auxing aux $
          fmap (subExpsRes . pure) . letSubExp "v"
            =<< eIndex arr (map toExp slice')
      insertIrregularM ns flags offsets (distResTag res) elems Dense env

transformFlatIndex ::
  TrCtx ->
  VName ->
  FlatSlice SubExp ->
  FlattenM DistEnv
transformFlatIndex :: TrCtx -> VName -> FlatSlice SubExp -> FlattenM DistEnv
transformFlatIndex (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) VName
arr FlatSlice SubExp
flat_slice
  | DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> FlatSlice SubExp -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) FlatSlice SubExp
flat_slice) = do
      arr_t <- DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
arr
      -- arr should be 1D
      let [n] = arrayDims arr_t
      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)
      arr_flat_size <- letSubExp "arr_flat_size" =<< toExp (pe64 num_segments * pe64 n)
      let arr_lift_shape = [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
arr_t
          arr_flat_shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
arr_flat_size]
      arr' <-
        liftSubExpRegular
          lvl
          segments
          inps
          env
          arr_lift_shape
          (Var arr)
      arr'_flat <-
        letExp (baseName arr <> "_reshaped") $ BasicOp $ Reshape arr' $ reshapeAll arr_lift_shape arr_flat_shape
      let FlatSlice off dims = flat_slice
          flat_slice' = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall d. d -> [FlatDimIndex d] -> FlatSlice d
FlatSlice SubExp
off (SubExp -> SubExp -> FlatDimIndex SubExp
forall d. d -> d -> FlatDimIndex d
FlatDimIndex SubExp
num_segments SubExp
n FlatDimIndex SubExp
-> [FlatDimIndex SubExp] -> [FlatDimIndex SubExp]
forall a. a -> [a] -> [a]
: [FlatDimIndex SubExp]
dims)
      out_flat_updated <-
        certifying (distCerts inps aux env) . letExp "flat_index_reg" . BasicOp $
          FlatIndex arr'_flat flat_slice'
      out_updated <-
        letExp "flat_index_reg_reshaped" $
          BasicOp $
            Reshape out_flat_updated $
              reshapeAll arr_flat_shape arr_lift_shape
      pure $ insertRegulars [distResTag res] [out_updated] env
  | DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) (FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice)) = do
      let space :: [SubExp]
space = Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ([SubExp] -> Shape
segmentsShape [SubExp]
segments) [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice
      v <-
        Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flat_index_reg_gather"
          (Exp GPU -> FlattenM VName)
-> ((NonEmpty SubExp -> FlattenM Result) -> FlattenM (Exp GPU))
-> (NonEmpty SubExp -> FlattenM Result)
-> FlattenM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Exp GPU -> FlattenM (Exp GPU)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Exp rep -> m (Exp rep)
renameExp
          (Exp GPU -> FlattenM (Exp GPU))
-> ((NonEmpty SubExp -> FlattenM Result) -> FlattenM (Exp GPU))
-> (NonEmpty SubExp -> FlattenM Result)
-> FlattenM (Exp GPU)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel
-> NonEmpty SubExp
-> (NonEmpty 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] -> NonEmpty SubExp
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList [SubExp]
space)
          ((NonEmpty SubExp -> FlattenM Result) -> FlattenM VName)
-> (NonEmpty SubExp -> FlattenM Result) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ \NonEmpty SubExp
is -> do
            let ([SubExp]
seg_is, [SubExp]
in_is) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
segmentsRank [SubExp]
segments) (NonEmpty SubExp -> [SubExp]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty SubExp
is)
            [SubExp] -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()
readInputs [SubExp]
segments DistEnv
env [SubExp]
seg_is DistInputs
inps
            let flat_slice' :: FlatSlice (TPrimExp Int64 VName)
flat_slice'@(FlatSlice TPrimExp Int64 VName
flat_offset [FlatDimIndex (TPrimExp Int64 VName)]
_) = (SubExp -> TPrimExp Int64 VName)
-> FlatSlice SubExp -> FlatSlice (TPrimExp Int64 VName)
forall a b. (a -> b) -> FlatSlice a -> FlatSlice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 FlatSlice SubExp
flat_slice
                flat_i :: TPrimExp Int64 VName
flat_i =
                  TPrimExp Int64 VName
flat_offset
                    TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((TPrimExp Int64 VName
 -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
(*) ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
in_is) (FlatSlice (TPrimExp Int64 VName) -> [TPrimExp Int64 VName]
forall d. FlatSlice d -> [d]
flatSliceStrides FlatSlice (TPrimExp Int64 VName)
flat_slice'))
            StmAux () -> FlattenM Result -> FlattenM Result
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
aux (FlattenM Result -> FlattenM Result)
-> FlattenM Result -> FlattenM Result
forall a b. (a -> b) -> a -> b
$
              (SubExp -> Result) -> FlattenM SubExp -> FlattenM Result
forall a b. (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([SubExp] -> Result
subExpsRes ([SubExp] -> Result) -> (SubExp -> [SubExp]) -> SubExp -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure) (FlattenM SubExp -> FlattenM Result)
-> (Exp GPU -> FlattenM SubExp) -> Exp GPU -> FlattenM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"v"
                (Exp GPU -> FlattenM Result)
-> FlattenM (Exp GPU) -> FlattenM Result
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
flat_i]
      pure $ insertRegulars [distResTag res] [v] env
  | Bool
otherwise = do
      -- Maximally nonuniform case.
      num_segments <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"num_segments" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp ([SubExp] -> TPrimExp Int64 VName
segmentCount [SubExp]
segments)
      ns <- letExp "slice_sizes" <=< renameExp <=< segMap lvl (MkSolo num_segments) $ \(MkSolo SubExp
segment) -> do
        segment_is <- [SubExp] -> SubExp -> FlattenM [SubExp]
segmentCoordsFromFlat [SubExp]
segments SubExp
segment
        slice_ns <- mapM (readInput segments env segment_is inps) $ flatSliceDims flat_slice
        fmap varsRes . letTupExp "n" <=< toExp $ product $ map pe64 slice_ns
      (_n, offsets, m) <- exScanAndSum lvl ns
      (_, _, repiota_D) <- doRepIota lvl ns
      flags <- genFlags lvl m offsets
      elems <- letExp "flat_index_irreg_elems" <=< renameExp <=< segMap lvl (NE.singleton m) $ \NonEmpty SubExp
is -> 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
repiota_D (NonEmpty (FlattenM (Exp (Rep FlattenM)))
-> [FlattenM (Exp (Rep FlattenM))]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (NonEmpty (FlattenM (Exp (Rep FlattenM)))
 -> [FlattenM (Exp (Rep FlattenM))])
-> NonEmpty (FlattenM (Exp (Rep FlattenM)))
-> [FlattenM (Exp (Rep FlattenM))]
forall a b. (a -> b) -> a -> b
$ (SubExp -> FlattenM (Exp (Rep FlattenM)))
-> NonEmpty SubExp -> NonEmpty (FlattenM (Exp (Rep FlattenM)))
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp NonEmpty SubExp
is)
        segment_start <- letSubExp "segment_start" =<< eIndex offsets [eSubExp segment]
        segment_is <- segmentCoordsFromFlat segments segment
        readInputs segments env segment_is inps
        let flat_slice'@(FlatSlice flat_offset _) = fmap pe64 flat_slice
            local_flat = SubExp -> TPrimExp Int64 VName
pe64 (NonEmpty SubExp -> SubExp
forall a. NonEmpty a -> a
NE.head NonEmpty SubExp
is) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- SubExp -> TPrimExp Int64 VName
pe64 SubExp
segment_start
            local_is = [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex (FlatSlice (TPrimExp Int64 VName) -> [TPrimExp Int64 VName]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice (TPrimExp Int64 VName)
flat_slice') TPrimExp Int64 VName
local_flat
            flat_i = TPrimExp Int64 VName
flat_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((TPrimExp Int64 VName
 -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
(*) [TPrimExp Int64 VName]
local_is (FlatSlice (TPrimExp Int64 VName) -> [TPrimExp Int64 VName]
forall d. FlatSlice d -> [d]
flatSliceStrides FlatSlice (TPrimExp Int64 VName)
flat_slice'))
        auxing aux $
          fmap (subExpsRes . pure) . letSubExp "v"
            =<< eIndex arr [toExp flat_i]
      insertIrregularM ns flags offsets (distResTag res) elems Dense env

transformIota ::
  TrCtx ->
  SubExp ->
  SubExp ->
  SubExp ->
  IntType ->
  FlattenM DistEnv
transformIota :: TrCtx -> SubExp -> SubExp -> SubExp -> IntType -> FlattenM DistEnv
transformIota (TrCtx SegLevel
_lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) SubExp
n SubExp
x SubExp
s IntType
it
  | DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
n),
    Bool -> Bool
not (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
x),
    Bool -> Bool
not (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
s) = do
      iota_row <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"iota_reg_row" (Exp (Rep FlattenM) -> FlattenM VName)
-> Exp (Rep FlattenM) -> FlattenM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
Iota SubExp
n SubExp
x SubExp
s IntType
it
      v <-
        letExp "iota_reg" $
          BasicOp $
            Replicate (segmentsShape segments) (Var iota_row)
      pure $ insertRegulars [distResTag res] [v] env
transformIota (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) SubExp
n (Constant PrimValue
x) (Constant PrimValue
s) IntType
Int64
  | PrimValue -> Bool
zeroIsh PrimValue
x,
    PrimValue -> Bool
oneIsh PrimValue
s = do
      ns <- SegLevel
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps SubExp
n
      (flags, offsets, elems) <- certifying (distCerts inps aux env) $ doSegIota lvl ns
      insertIrregularM ns flags offsets (distResTag res) elems Dense env
transformIota (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) SubExp
n SubExp
x SubExp
s IntType
it = do
  ns <- SegLevel
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps SubExp
n
  xs <- dataArr lvl segments env inps x
  ss <- dataArr lvl segments env inps s
  (res_F, res_O, res_D) <- certifying (distCerts inps aux env) $ doSegIota lvl ns
  (_, _, repiota_D) <- doRepIota lvl ns
  m <- arraySize 0 <$> lookupType res_D
  res_D' <- letExp "iota_D_fixed" <=< segMap lvl (MkSolo m) $ \(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
repiota_D [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    v' <- letSubExp "v" =<< eIndex res_D [eSubExp i]
    x' <- letSubExp "x" =<< eIndex xs [eSubExp segment]
    s' <- letSubExp "s" =<< eIndex ss [eSubExp segment]
    fmap (subExpsRes . pure) . letSubExp "v" <=< toExp $
      primExpFromSubExp (IntType it) x'
        ~+~ sExt it (untyped (pe64 v'))
        ~*~ primExpFromSubExp (IntType it) s'
  insertIrregularM ns res_F res_O (distResTag res) res_D' Dense env

transformConcat ::
  TrCtx ->
  Int ->
  NE.NonEmpty VName ->
  SubExp ->
  FlattenM DistEnv
transformConcat :: TrCtx -> Int -> NonEmpty VName -> SubExp -> FlattenM DistEnv
transformConcat (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) Int
d NonEmpty VName
arr SubExp
shp = do
  arr_ts <- (VName -> FlattenM Type) -> [VName] -> FlattenM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps) (NonEmpty VName -> [VName]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty VName
arr)
  let inputShapeUniform TypeBase Shape m
t =
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) (TypeBase Shape m -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims TypeBase Shape m
t)
  if isRegularDistResult res
    && not (isVariant inps shp)
    && all inputShapeUniform arr_ts
    then do
      --  Unifrom Concat
      arrs_lifted <-
        forM (zip (NE.toList arr) arr_ts) $ \(VName
v, Type
t) -> do
          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
          SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
expectedShape (VName -> SubExp
Var VName
v)
      v' <-
        letExp "concat_reg" $
          BasicOp $
            Concat
              (segmentsRank segments + d)
              (NE.fromList arrs_lifted)
              shp

      pure $ insertRegulars [distResTag res] [v'] env
    else do
      ns <- dataArr lvl segments env inps shp
      reparr <- mapM (getIrregRep lvl segments env inps) (NE.toList arr)
      rep' <- case d of
        Int
0 -> SegLevel
-> [SubExp]
-> DistEnv
-> VName
-> [IrregularRep]
-> FlattenM IrregularRep
concatIrreg SegLevel
lvl [SubExp]
segments DistEnv
env VName
ns [IrregularRep]
reparr
        Int
d' -> do
          SegLevel
-> [SubExp]
-> DistEnv
-> VName
-> [IrregularRep]
-> [Type]
-> DistInputs
-> Int
-> FlattenM IrregularRep
concatIrregAlongDim SegLevel
lvl [SubExp]
segments DistEnv
env VName
ns [IrregularRep]
reparr [Type]
arr_ts DistInputs
inps Int
d'
      insertRepM (distResTag res) (Irregular rep') env

transformReplicate ::
  TrCtx ->
  Shape ->
  SubExp ->
  FlattenM DistEnv
--  Uniform Replicate
transformReplicate :: TrCtx -> Shape -> SubExp -> FlattenM DistEnv
transformReplicate (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) (Shape [SubExp]
dims) SubExp
se
  | DistResult -> Bool
isRegularDistResult DistResult
res = do
      t <- DistInputs -> SubExp -> FlattenM Type
subExpInputType DistInputs
inps SubExp
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
t
      lifted <- liftSubExpRegular lvl segments inps env expectedShape se
      v_rep <- replicateForDims segments (Shape dims) lifted
      pure $ insertRegulars [distResTag res] [v_rep] env
transformReplicate (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) (Shape [SubExp
n]) (Var VName
v) = do
  ns <- SegLevel
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps SubExp
n
  rep <- getIrregRep lvl segments env inps v
  rep' <- replicateIrreg lvl segments env ns (baseName v) rep
  insertRepM (distResTag res) (Irregular rep') env
transformReplicate (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) (Shape [SubExp
n]) (Constant PrimValue
v) = do
  ns <- SegLevel
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps SubExp
n
  (res_F, res_O, res_D) <-
    certifying (distCerts inps aux env) $ doSegIota lvl ns
  w <- arraySize 0 <$> lookupType res_D
  res_D' <- letExp "rep_const" $ BasicOp $ Replicate (Shape [w]) (Constant v)
  insertIrregularM ns res_F res_O (distResTag res) res_D' Dense env
transformReplicate (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) (Shape [SubExp]
dims) (Constant PrimValue
v) = do
  dim_arrs <- (SubExp -> FlattenM VName) -> [SubExp] -> FlattenM [VName]
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
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps) [SubExp]
dims
  seg_number <- arraySize 0 <$> lookupType (head dim_arrs)
  mul_dims <- letExp "mul_dims" <=< segMap lvl (MkSolo seg_number) $ \(MkSolo SubExp
i) -> do
    vals <- (VName -> FlattenM SubExp) -> [VName] -> 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 (\VName
dim_arr -> Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"dim_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
dim_arr [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]) [VName]
dim_arrs
    n <- letSubExp "n" <=< toExp $ product $ map pe64 vals
    pure [subExpRes n]
  (res_F, res_O, res_D) <-
    certifying (distCerts inps aux env) $ doSegIota lvl mul_dims
  w <- arraySize 0 <$> lookupType res_D
  res_D' <- letExp "rep_const" $ BasicOp $ Replicate (Shape [w]) (Constant v)
  insertIrregularM mul_dims res_F res_O (distResTag res) res_D' Dense env
transformReplicate (TrCtx SegLevel
_lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) (Shape []) (Var VName
v) =
  case VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps of
    Just (DistInputFree VName
v' Type
_) -> 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
"_copy") (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 Shape
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
v')
      pure $ insertRegulars [distResTag res] [v''] env
    Just (DistInput ResTag
rt Type
_) ->
      case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
        Irregular IrregularRep
r -> do
          let name :: Name
name = VName -> Name
baseName (IrregularRep -> VName
irregularD IrregularRep
r) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_copy"
          elems_copy <-
            Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
name (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 Shape
forall a. Monoid a => a
mempty (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ IrregularRep -> VName
irregularD IrregularRep
r)
          let rep = IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep) -> IrregularRep -> ResRep
forall a b. (a -> b) -> a -> b
$ IrregularRep
r {irregularD = elems_copy}
          insertRepM (distResTag res) rep env
        Regular VName
v' -> 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
"_copy") (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 Shape
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
v')
          pure $ insertRegulars [distResTag res] [v''] env
    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
"_copy_free") (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)
      pure $ insertRegulars [distResTag res] [v'] env
transformReplicate (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) (Shape [SubExp]
dims) (Var VName
v) = do
  dim_arrs <- (SubExp -> FlattenM VName) -> [SubExp] -> FlattenM [VName]
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
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps) [SubExp]
dims
  seg_number <- arraySize 0 <$> lookupType (head dim_arrs)
  mul_dims <- letExp "mul_dims" <=< segMap lvl (MkSolo seg_number) $ \(MkSolo SubExp
i) -> do
    vals <- (VName -> FlattenM SubExp) -> [VName] -> 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 (\VName
dim_arr -> Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"dim_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
dim_arr [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]) [VName]
dim_arrs
    n <- letSubExp "n" <=< toExp $ product $ map pe64 vals
    pure [subExpRes n]
  rep <- getIrregRep lvl segments env inps v
  rep' <- replicateIrreg lvl segments env mul_dims (baseName v) rep
  insertRepM (distResTag res) (Irregular rep') env

transformManifest ::
  TrCtx ->
  VName ->
  [Int] ->
  FlattenM DistEnv
transformManifest :: TrCtx -> VName -> [Int] -> FlattenM DistEnv
transformManifest (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) VName
v [Int]
perm
  | DistResult -> Bool
isRegularDistResult DistResult
res = do
      t <- DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
v
      v_lifted <-
        liftSubExpRegular
          lvl
          segments
          inps
          env
          (segmentsShape segments <> arrayShape t)
          (Var v)
      let segment_rank = [SubExp] -> Int
segmentsRank [SubExp]
segments
      v_manifest <- letExp (baseName v <> "_manifest") . BasicOp $ Manifest v_lifted ([0 .. segment_rank - 1] ++ map (+ segment_rank) perm)
      pure $ insertRegulars [distResTag res] [v_manifest] env
  | Bool
otherwise = do
      irreg <- SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
v
      irreg_dense <- ensureDenseIrregular lvl (baseName v <> "_manifest") irreg
      elems_copy <-
        letExp (baseName (irregularD irreg_dense) <> "_manifest") . BasicOp $
          Replicate mempty (Var $ irregularD irreg_dense)
      insertRepM
        (distResTag res)
        (Irregular $ irreg_dense {irregularD = elems_copy})
        env

transformUpdate ::
  TrCtx ->
  Safety ->
  VName ->
  Slice SubExp ->
  SubExp ->
  FlattenM DistEnv
transformUpdate :: TrCtx
-> Safety -> VName -> Slice SubExp -> SubExp -> FlattenM DistEnv
transformUpdate (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) Safety
safety VName
as Slice SubExp
slice SubExp
se
  -- Uniform Update
  | Just Type
as_t <- DistInput -> Type
distInputType (DistInput -> Type) -> Maybe DistInput -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
as DistInputs
inps,
    DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> Slice SubExp -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) Slice SubExp
slice) = do
      as' <-
        SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular
          SegLevel
lvl
          [SubExp]
segments
          DistInputs
inps
          DistEnv
env
          ([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
as_t)
          (VName -> SubExp
Var VName
as)
      se' <-
        liftSubExpRegular
          lvl
          segments
          inps
          env
          (segmentsShape segments <> sliceShape slice)
          se
      let segmentSlice = (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim ([SubExp] -> [DimIndex SubExp])
-> ([SubExp] -> [SubExp]) -> [SubExp] -> [DimIndex SubExp]
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
      v <-
        certifying (distCerts inps aux env) . letExp "update_reg" . BasicOp $
          Update safety as' (Slice $ segmentSlice segments <> unSlice slice) (Var se')
      pure $ insertRegulars [distResTag res] [v] env
  | Just Type
as_t <- DistInput -> Type
distInputType (DistInput -> Type) -> Maybe DistInput -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
as DistInputs
inps,
    DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) (Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice)) = do
      let as_lift_shape :: Shape
as_lift_shape = [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
as_t
      as' <- SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
as_lift_shape (VName -> SubExp
Var VName
as)
      let update_dims = [SubExp]
segments [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
      updated <-
        certifying (distCerts inps aux env)
          . letExp "update_reg_scatter"
          <=< renameExp
          <=< genScatterND lvl as' update_dims
          $ \[SubExp]
is -> do
            let ([SubExp]
seg_is, [SubExp]
in_is) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
segmentsRank [SubExp]
segments) ([SubExp] -> [SubExp]
forall a. [a] -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [SubExp]
is)
            [SubExp] -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()
readInputs [SubExp]
segments DistEnv
env [SubExp]
seg_is (DistInputs -> FlattenM ()) -> DistInputs -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ ((VName, DistInput) -> Bool) -> DistInputs -> DistInputs
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
as) (VName -> Bool)
-> ((VName, DistInput) -> VName) -> (VName, DistInput) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, DistInput) -> VName
forall a b. (a, b) -> a
fst) DistInputs
inps
            let slice' :: [TPrimExp Int64 VName]
slice' = Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice ((SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice) ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
in_is)
            -- Value to write
            v' <- case SubExp
se of
              Constant PrimValue
c -> 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
c
              Var VName
se_v -> Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"v" (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
se_v ((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 a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp [SubExp]
in_is)
            -- Index to write `v'` at
            in_is' <- mapM (letSubExp "i" <=< toExp) slice'
            let is' = [SubExp]
seg_is [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> [SubExp]
in_is'
            pure (is', v')
      pure $ insertRegulars [distResTag res] [updated] env
  | Just Type
as_t <- DistInput -> Type
distInputType (DistInput -> Type) -> Maybe DistInput -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
as DistInputs
inps = do
      num_segments <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"num_segments" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp ([SubExp] -> TPrimExp Int64 VName
segmentCount [SubExp]
segments)
      ns <- letExp "slice_sizes"
        <=< renameExp
        <=< segMap lvl (MkSolo num_segments)
        $ \(MkSolo SubExp
seg_i) -> do
          seg_is <- [SubExp] -> SubExp -> FlattenM [SubExp]
segmentCoordsFromFlat [SubExp]
segments SubExp
seg_i
          readInputs segments env seg_is $
            filter ((`elem` sliceDims slice) . Var . fst) inps
          slice_dims <- mapM (readInput segments env seg_is inps) $ sliceDims slice
          n <- letSubExp "n" <=< toExp $ product $ map pe64 slice_dims
          pure [subExpRes n]
      -- Irregular representation of `as`
      as_rep <- getIrregRep lvl segments env inps as
      IrregularRep shape flags offsets elems _ <-
        ensureDenseIrregular lvl (baseName as <> "_update") as_rep
      -- Inner indices (1 and 2) of `ns`
      (_, _, ii1_vss) <- doRepIota lvl ns
      (_, _, ii2_vss) <- certifying (distCerts inps aux env) $ doSegIota lvl ns
      -- Number of updates to perform
      m <- arraySize 0 <$> lookupType ii2_vss
      elems' <- letExp "elems_scatter" <=< renameExp <=< genScatter lvl elems m $ \SubExp
gid -> 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
ii1_vss [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gid]
        in_seg_i <- letSubExp "in_seg_i" =<< eIndex ii2_vss [eSubExp gid]
        seg_is <- segmentCoordsFromFlat segments seg_i
        readInputs segments env seg_is $ filter ((/= as) . fst) inps
        as_dims <- readTypeDims segments env seg_is inps as_t
        slice_dims <- mapM (readInput segments env seg_is inps) $ sliceDims slice
        case se of
          Var VName
v -> do
            let in_seg_is :: [TPrimExp Int64 VName]
in_seg_is =
                  [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
slice_dims) (SubExp -> TPrimExp Int64 VName
pe64 SubExp
in_seg_i)
                slice' :: Slice (TPrimExp Int64 VName)
slice' = (SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice
                flat_i :: TPrimExp Int64 VName
flat_i =
                  [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall num. IntegralExp num => [num] -> [num] -> num
flattenIndex
                    ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
as_dims)
                    (Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice Slice (TPrimExp Int64 VName)
slice' [TPrimExp Int64 VName]
in_seg_is)
            -- Value to write
            v' <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"v" (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 ((TPrimExp Int64 VName -> FlattenM (Exp GPU))
-> [TPrimExp Int64 VName] -> [FlattenM (Exp GPU)]
forall a b. (a -> b) -> [a] -> [b]
map 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]
in_seg_is)
            o' <- letSubExp "o" =<< eIndex offsets [eSubExp seg_i]
            -- Index to write `v'` at
            i <- letExp "i" =<< toExp (pe64 o' + flat_i)
            pure (i, v')
          Constant PrimValue
c -> do
            let slice' :: Slice (TPrimExp Int64 VName)
slice' = (SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice
                flat_i :: TPrimExp Int64 VName
flat_i = [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall num. IntegralExp num => [num] -> [num] -> num
flattenIndex ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
as_dims) (Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice Slice (TPrimExp Int64 VName)
slice' [])
            o' <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"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
offsets [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
seg_i]
            i <- letExp "i" =<< toExp (pe64 o' + flat_i)
            pure (i, Constant c)
      insertIrregularM shape flags offsets (distResTag res) elems' Dense env
  | Bool
otherwise =
      String -> FlattenM DistEnv
forall a. HasCallStack => String -> a
error String
"Flattening update: destination is not input."

transformFlatUpdate ::
  TrCtx ->
  VName ->
  FlatSlice SubExp ->
  VName ->
  FlattenM DistEnv
transformFlatUpdate :: TrCtx -> VName -> FlatSlice SubExp -> VName -> FlattenM DistEnv
transformFlatUpdate (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) VName
as FlatSlice SubExp
flat_slice VName
v
  -- Uniform Flat Update
  | Just Type
as_t <- DistInput -> Type
distInputType (DistInput -> Type) -> Maybe DistInput -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
as DistInputs
inps,
    DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> FlatSlice SubExp -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) FlatSlice SubExp
flat_slice) = do
      -- as should be 1D
      let [SubExp
n] = Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
as_t
      num_segments <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"num_segments" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp ([SubExp] -> TPrimExp Int64 VName
segmentCount [SubExp]
segments)
      as_flat_size <- letSubExp "as_flat_size" =<< toExp (pe64 num_segments * pe64 n)

      let se_shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice
          as_lift_shape = [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
as_t
          se_lift_shape = [SubExp] -> Shape
segmentsShape [SubExp]
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
se_shape
          se_flat_shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
num_segments] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
se_shape
          as_flat_shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
as_flat_size]
      as' <-
        liftSubExpRegular
          lvl
          segments
          inps
          env
          as_lift_shape
          (Var as)
      v' <-
        liftSubExpRegular
          lvl
          segments
          inps
          env
          se_lift_shape
          (Var v)
      as_flat <-
        letExp (baseName as <> "_reshaped") $ BasicOp $ Reshape as' $ reshapeAll as_lift_shape as_flat_shape
      v_flat <-
        letExp (baseName v <> "_reshaped") $ BasicOp $ Reshape v' $ reshapeAll se_lift_shape se_flat_shape
      let FlatSlice off dims = flat_slice
          flat_slice' = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall d. d -> [FlatDimIndex d] -> FlatSlice d
FlatSlice SubExp
off (SubExp -> SubExp -> FlatDimIndex SubExp
forall d. d -> d -> FlatDimIndex d
FlatDimIndex SubExp
num_segments SubExp
n FlatDimIndex SubExp
-> [FlatDimIndex SubExp] -> [FlatDimIndex SubExp]
forall a. a -> [a] -> [a]
: [FlatDimIndex SubExp]
dims)
      out_flat_updated <-
        certifying (distCerts inps aux env) . letExp "flat_update_reg" . BasicOp $
          FlatUpdate as_flat flat_slice' v_flat
      out_updated <-
        letExp "flat_update_reg_reshaped" $
          BasicOp $
            Reshape out_flat_updated $
              reshapeAll as_flat_shape as_lift_shape
      pure $ insertRegulars [distResTag res] [out_updated] env
  | Just Type
as_t <- DistInput -> Type
distInputType (DistInput -> Type) -> Maybe DistInput -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
as DistInputs
inps,
    DistResult -> Bool
isRegularDistResult DistResult
res,
    Bool -> Bool
not ((SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) (FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice)) = do
      -- as should be 1D
      let as_lift_shape :: Shape
as_lift_shape = [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
as_t
      as' <- SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
as_lift_shape (VName -> SubExp
Var VName
as)
      let update_dims = [SubExp]
segments [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice
      updated <-
        certifying (distCerts inps aux env)
          . letExp "flat_update_reg_scatter"
          <=< renameExp
          <=< genScatterND lvl as' update_dims
          $ \[SubExp]
is -> do
            let ([SubExp]
seg_is, [SubExp]
in_is) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
segmentsRank [SubExp]
segments) ([SubExp] -> [SubExp]
forall a. [a] -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [SubExp]
is)
            [SubExp] -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()
readInputs [SubExp]
segments DistEnv
env [SubExp]
seg_is (DistInputs -> FlattenM ()) -> DistInputs -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ ((VName, DistInput) -> Bool) -> DistInputs -> DistInputs
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
as) (VName -> Bool)
-> ((VName, DistInput) -> VName) -> (VName, DistInput) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, DistInput) -> VName
forall a b. (a, b) -> a
fst) DistInputs
inps
            let flat_slice' :: FlatSlice (TPrimExp Int64 VName)
flat_slice'@(FlatSlice TPrimExp Int64 VName
flat_offset [FlatDimIndex (TPrimExp Int64 VName)]
_) = (SubExp -> TPrimExp Int64 VName)
-> FlatSlice SubExp -> FlatSlice (TPrimExp Int64 VName)
forall a b. (a -> b) -> FlatSlice a -> FlatSlice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 FlatSlice SubExp
flat_slice
                flat_i :: TPrimExp Int64 VName
flat_i =
                  TPrimExp Int64 VName
flat_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((TPrimExp Int64 VName
 -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
(*) ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
in_is) (FlatSlice (TPrimExp Int64 VName) -> [TPrimExp Int64 VName]
forall d. FlatSlice d -> [d]
flatSliceStrides FlatSlice (TPrimExp Int64 VName)
flat_slice'))
            -- Value to write
            v' <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"v" (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 ((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 a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp [SubExp]
in_is)
            -- Index to write `v'` at
            i <- letSubExp "i" =<< toExp flat_i
            pure (seg_is <> [i], v')
      pure $ insertRegulars [distResTag res] [updated] env
  | Just Type
_ <- DistInput -> Type
distInputType (DistInput -> Type) -> Maybe DistInput -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
as DistInputs
inps = do
      num_segments <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"num_segments" (Exp GPU -> FlattenM SubExp)
-> FlattenM (Exp GPU) -> FlattenM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp ([SubExp] -> TPrimExp Int64 VName
segmentCount [SubExp]
segments)
      ns <- letExp "slice_sizes"
        <=< renameExp
        <=< segMap lvl (MkSolo num_segments)
        $ \(MkSolo SubExp
seg_i) -> do
          seg_is <- [SubExp] -> SubExp -> FlattenM [SubExp]
segmentCoordsFromFlat [SubExp]
segments SubExp
seg_i
          readInputs segments env seg_is $
            filter ((`elem` flatSliceDims flat_slice) . Var . fst) inps
          slice_dims <- mapM (readInput segments env seg_is inps) $ flatSliceDims flat_slice
          n <- letSubExp "n" <=< toExp $ product $ map pe64 slice_dims
          pure [subExpRes n]
      -- Irregular representation of `as`
      as_rep <- getIrregRep lvl segments env inps as
      IrregularRep shape flags offsets elems _ <-
        ensureDenseIrregular lvl (baseName as <> "_update") as_rep
      -- Inner indices (1 and 2) of `ns`
      (_, _, ii1_vss) <- doRepIota lvl ns
      (_, _, ii2_vss) <- certifying (distCerts inps aux env) $ doSegIota lvl ns
      -- Number of updates to perform
      m <- arraySize 0 <$> lookupType ii2_vss
      elems' <- letExp "elems_scatter" <=< renameExp <=< genScatter lvl elems m $ \SubExp
gid -> 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
ii1_vss [SubExp -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gid]
        in_seg_i <- letSubExp "in_seg_i" =<< eIndex ii2_vss [eSubExp gid]
        seg_is <- segmentCoordsFromFlat segments seg_i
        readInputs segments env seg_is $ filter ((/= as) . fst) inps
        let slice_dims = FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
flat_slice
            flat_stride = FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceStrides FlatSlice SubExp
flat_slice
            (FlatSlice flat_offset _) = fmap pe64 flat_slice
            in_seg_is =
              [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
slice_dims) (SubExp -> TPrimExp Int64 VName
pe64 SubExp
in_seg_i)
            flat_i = TPrimExp Int64 VName
flat_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((TPrimExp Int64 VName
 -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
(*) [TPrimExp Int64 VName]
in_seg_is ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
flat_stride))
        -- Value to write
        v' <- letSubExp "v" =<< eIndex v (map toExp in_seg_is)
        o' <- letSubExp "o" =<< eIndex offsets [eSubExp seg_i]
        -- Index to write `v'` at
        i <- letExp "i" =<< toExp (pe64 o' + flat_i)
        pure (i, v')
      insertIrregularM shape flags offsets (distResTag res) elems' Dense env
  | Bool
otherwise =
      String -> FlattenM DistEnv
forall a. HasCallStack => String -> a
error String
"Flattening update: destination is not input."

transformRearrange ::
  TrCtx ->
  VName ->
  [Int] ->
  FlattenM DistEnv
transformRearrange :: TrCtx -> VName -> [Int] -> FlattenM DistEnv
transformRearrange (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux) VName
v [Int]
perm
  | DistResult -> Bool
isRegularDistResult DistResult
res = do
      t <- DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
v
      v_lifted <-
        liftSubExpRegular
          lvl
          segments
          inps
          env
          (segmentsShape segments <> arrayShape t)
          (Var v)
      let segment_rank = [SubExp] -> Int
segmentsRank [SubExp]
segments
      v_rearrange <- letExp (baseName v <> "_tr") . BasicOp $ Rearrange v_lifted ([0 .. segment_rank - 1] ++ map (+ segment_rank) perm)
      pure $ insertRegulars [distResTag res] [v_rearrange] env
  | Bool
otherwise = do
      irreg <- SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps VName
v
      t <- lookupInputType inps v
      rep' <-
        certifying (distCerts inps aux env) $
          rearrangeIrreg lvl segments env inps t perm irreg
      insertRepM (distResTag res) (Irregular rep') env

transformScratch ::
  TrCtx ->
  PrimType ->
  [SubExp] ->
  FlattenM DistEnv
transformScratch :: TrCtx -> PrimType -> [SubExp] -> FlattenM DistEnv
transformScratch (TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
_aux) PrimType
pt [SubExp]
dims
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) [SubExp]
dims = do
      -- All dims are invariant result is regular across segments.
      v' <-
        Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"scratch" (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
$
          PrimType -> [SubExp] -> BasicOp
Scratch PrimType
pt (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ([SubExp] -> Shape
segmentsShape [SubExp]
segments) [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
dims)
      pure $ insertRegulars [distResTag res] [v'] env
  | [SubExp
n] <- [SubExp]
dims = do
      ns <- SegLevel
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps SubExp
n
      (_n, offsets, m) <- exScanAndSum lvl ns
      flags <- genFlags lvl m offsets
      res_D <- letExp "scratch_D" $ BasicOp $ Scratch pt [m]
      insertIrregularM ns flags offsets (distResTag res) res_D Dense env
  | Bool
otherwise = do
      dim_arrs <- (SubExp -> FlattenM VName) -> [SubExp] -> FlattenM [VName]
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
-> [SubExp] -> DistEnv -> DistInputs -> SubExp -> FlattenM VName
dataArr SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps) [SubExp]
dims
      w <- arraySize 0 <$> lookupType (head dim_arrs)
      ns <- letExp "scratch_sizes" <=< segMap lvl (MkSolo w) $ \(MkSolo SubExp
i) -> do
        vals <- (VName -> FlattenM SubExp) -> [VName] -> 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 (\VName
arr -> Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"d" (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 (Rep FlattenM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]) [VName]
dim_arrs
        n <- letSubExp "n" <=< toExp $ product $ map pe64 vals
        pure [subExpRes n]
      (_n, offsets, m) <- exScanAndSum lvl ns
      flags <- genFlags lvl m offsets
      res_D <- letExp "scratch_D" $ BasicOp $ Scratch pt [m]
      insertIrregularM ns flags offsets (distResTag res) res_D Dense env

flattenBasicOp ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  ( DistInputs,
    DistResult,
    PatElem Type,
    StmAux (),
    BasicOp
  ) ->
  FlattenM DistEnv
flattenBasicOp :: FlattenOps
-> [SubExp]
-> DistEnv
-> (DistInputs, DistResult, PatElem Type, StmAux (), BasicOp)
-> FlattenM DistEnv
flattenBasicOp FlattenOps
ops [SubExp]
segments DistEnv
env (DistInputs
inps, DistResult
res, PatElem Type
pe, StmAux ()
aux, BasicOp
e) =
  case BasicOp
e of
    BinOp {} -> FlattenM DistEnv
scalarCase
    CmpOp {} -> FlattenM DistEnv
scalarCase
    ConvOp {} -> FlattenM DistEnv
scalarCase
    UnOp {} -> FlattenM DistEnv
scalarCase
    UserParam Name
_ SubExp
_ -> FlattenM DistEnv
scalarCase -- These are always of type i64.
    Assert {} -> FlattenM DistEnv
scalarCase
    ArrayLit [SubExp]
vs Type
row_type -> TrCtx -> [SubExp] -> Type -> FlattenM DistEnv
transformArrayLit TrCtx
ctx [SubExp]
vs Type
row_type
    ArrayVal [PrimValue]
vs PrimType
row_type -> TrCtx -> [PrimValue] -> PrimType -> FlattenM DistEnv
transformArrayVal TrCtx
ctx [PrimValue]
vs PrimType
row_type
    Opaque OpaqueOp
op SubExp
se -> (SubExp -> BasicOp) -> SubExp -> FlattenM DistEnv
passThrough (OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
op) SubExp
se
    Reshape VName
arr NewShape SubExp
reshape -> TrCtx -> VName -> NewShape SubExp -> FlattenM DistEnv
transformReshape TrCtx
ctx VName
arr NewShape SubExp
reshape
    Index VName
arr Slice SubExp
slice -> TrCtx -> VName -> Slice SubExp -> FlattenM DistEnv
transformIndex TrCtx
ctx VName
arr Slice SubExp
slice
    FlatIndex VName
arr FlatSlice SubExp
flat_slice -> TrCtx -> VName -> FlatSlice SubExp -> FlattenM DistEnv
transformFlatIndex TrCtx
ctx VName
arr FlatSlice SubExp
flat_slice
    Iota SubExp
n SubExp
x SubExp
s IntType
it -> TrCtx -> SubExp -> SubExp -> SubExp -> IntType -> FlattenM DistEnv
transformIota TrCtx
ctx SubExp
n SubExp
x SubExp
s IntType
it
    Concat Int
d NonEmpty VName
arr SubExp
shp -> TrCtx -> Int -> NonEmpty VName -> SubExp -> FlattenM DistEnv
transformConcat TrCtx
ctx Int
d NonEmpty VName
arr SubExp
shp
    Replicate Shape
shape SubExp
se -> TrCtx -> Shape -> SubExp -> FlattenM DistEnv
transformReplicate TrCtx
ctx Shape
shape SubExp
se
    Manifest VName
v [Int]
perm -> TrCtx -> VName -> [Int] -> FlattenM DistEnv
transformManifest TrCtx
ctx VName
v [Int]
perm
    Update Safety
safety VName
as Slice SubExp
slice SubExp
se -> TrCtx
-> Safety -> VName -> Slice SubExp -> SubExp -> FlattenM DistEnv
transformUpdate TrCtx
ctx Safety
safety VName
as Slice SubExp
slice SubExp
se
    FlatUpdate VName
as FlatSlice SubExp
flat_slice VName
v -> TrCtx -> VName -> FlatSlice SubExp -> VName -> FlattenM DistEnv
transformFlatUpdate TrCtx
ctx VName
as FlatSlice SubExp
flat_slice VName
v
    Rearrange VName
v [Int]
perm -> TrCtx -> VName -> [Int] -> FlattenM DistEnv
transformRearrange TrCtx
ctx VName
v [Int]
perm
    Scratch PrimType
pt [SubExp]
dims -> TrCtx -> PrimType -> [SubExp] -> FlattenM DistEnv
transformScratch TrCtx
ctx PrimType
pt [SubExp]
dims
    UpdateAcc {} ->
      -- TODO: handle nonuniform case, which is however rare, and also needs
      -- modifications to WithAcc. The only irregularity that is possible is in
      -- the values to be written.
      FlattenM DistEnv
scalarCase
    SubExp SubExp
se -> (SubExp -> BasicOp) -> SubExp -> FlattenM DistEnv
passThrough SubExp -> BasicOp
SubExp SubExp
se
  where
    ctx :: TrCtx
ctx = SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> DistResult
-> StmAux ()
-> TrCtx
TrCtx SegLevel
lvl [SubExp]
segments DistEnv
env DistInputs
inps DistResult
res StmAux ()
aux
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops
    scalarCase :: FlattenM DistEnv
scalarCase =
      FlattenOps
-> [SubExp]
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
flattenScalarStm FlattenOps
ops [SubExp]
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 ([PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem Type
pe]) StmAux ()
StmAux (ExpDec SOACS)
aux (BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp BasicOp
e)

    -- Distribute a BasicOp that merely passes its operand through, applying it
    -- to the representation of a distributed input (or falling back to the
    -- scalar case).
    passThrough :: (SubExp -> BasicOp) -> SubExp -> FlattenM DistEnv
passThrough SubExp -> BasicOp
mkOp SubExp
se
      | Var VName
v <- SubExp
se,
        Just (DistInput ResTag
rt_in Type
_) <- VName -> DistInputs -> Maybe DistInput
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v DistInputs
inps =
          case ResTag -> DistEnv -> ResRep
resVar ResTag
rt_in DistEnv
env of
            Regular VName
arr -> do
              arr' <-
                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)
-> (Exp GPU -> FlattenM VName) -> Exp GPU -> FlattenM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) -> Exp GPU -> FlattenM VName
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> Exp GPU
forall rep. BasicOp -> Exp rep
BasicOp (SubExp -> BasicOp
mkOp (VName -> SubExp
Var VName
arr))
              pure $ insertRegulars [distResTag res] [arr'] env
            Irregular IrregularRep
irreg -> do
              elems' <-
                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)
-> (Exp GPU -> FlattenM VName) -> Exp GPU -> FlattenM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) -> Exp GPU -> FlattenM VName
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> Exp GPU
forall rep. BasicOp -> Exp rep
BasicOp (SubExp -> BasicOp
mkOp (VName -> SubExp
Var (IrregularRep -> VName
irregularD IrregularRep
irreg)))
              insertRepM (distResTag res) (Irregular irreg {irregularD = elems'}) env
      | Bool
otherwise =
          FlattenM DistEnv
scalarCase