-- | Flattening of 'Match'.
module Futhark.Pass.Flatten.Match
  ( flattenMatch,
  )
where

import Control.Monad
import Data.Containers.ListUtils (nubOrd)
import Data.List qualified as L
import Data.Map qualified as M
import Data.Maybe
import Data.Set qualified as S
import Data.Tuple.Solo
import Futhark.IR.GPU
import Futhark.IR.SOACS
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Tools

-- Given the indices for which a branch is taken and its body,
-- distribute the statements of the body of that branch.
distributeBranch ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  VName ->
  Body SOACS ->
  M.Map VName ResRep ->
  FlattenM (DistInputs, DistEnv, DistStms)
distributeBranch :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> VName
-> Body SOACS
-> Map VName ResRep
-> FlattenM (DistInputs, DistEnv, DistStms)
distributeBranch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps VName
is Body SOACS
body Map VName ResRep
acc_reps = do
  let lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops
      free_in_body :: [VName]
free_in_body = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (SubExp -> Bool) -> (VName -> SubExp) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) (Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Body SOACS
body)
  scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  free_sizes <-
    foldMap freeIn <$> mapM (lookupInputType inps) free_in_body
  let free_variant_sizes = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (SubExp -> Bool) -> (VName -> SubExp) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) (Names -> [VName]
namesToList Names
free_sizes)
      free_size_vars = [VName] -> [VName]
forall a. Ord a => [a] -> [a]
nubOrd ([VName]
free_variant_sizes [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
free_in_body)
  (ts, vs, reps) <-
    unzip3 <$> mapM (splitInput lvl segments env inps is acc_reps) free_size_vars
  let inputs = do
        (v, t, i) <- [VName] -> [Type] -> [Int] -> [(VName, Type, Int)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [VName]
vs [Type]
ts [Int
0 ..]
        pure (v, DistInput (ResTag i) t)
  let env' = Map ResTag ResRep -> DistEnv
DistEnv (Map ResTag ResRep -> DistEnv) -> Map ResTag ResRep -> DistEnv
forall a b. (a -> b) -> a -> b
$ [(ResTag, ResRep)] -> Map ResTag ResRep
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(ResTag, ResRep)] -> Map ResTag ResRep)
-> [(ResTag, ResRep)] -> Map ResTag ResRep
forall a b. (a -> b) -> a -> b
$ [ResTag] -> [ResRep] -> [(ResTag, ResRep)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> ResTag) -> [Int] -> [ResTag]
forall a b. (a -> b) -> [a] -> [b]
map Int -> ResTag
ResTag [Int
0 ..]) [ResRep]
reps
  let (inputs', dstms) = distributeBodyWith ops scope segments inputs body
  pure (inputs', env', dstms)

-- Given a single result from each branch as well the *unlifted*
-- result type, merge the results of all branches into a single result.
mergeResult ::
  SegLevel ->
  Segments ->
  SubExp ->
  [VName] ->
  [ResRep] ->
  DistResult ->
  FlattenM ResRep
mergeResult :: SegLevel
-> Segments
-> SubExp
-> [VName]
-> [ResRep]
-> DistResult
-> FlattenM ResRep
mergeResult SegLevel
lvl Segments
segments SubExp
w [VName]
iss [ResRep]
branchesRep DistResult
dist_res
  -- Regular case
  | DistResult -> Bool
isRegularDistResult DistResult
dist_res = do
      let (DistType Segments
_ Rank
_ Type
resType) = DistResult -> DistType
distResType DistResult
dist_res
      if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
resType
        then do
          xs <- (ResRep -> FlattenM VName) -> [ResRep] -> 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 ResRep -> FlattenM VName
forall {f :: * -> *}. Applicative f => ResRep -> f VName
asRegular [ResRep]
branchesRep
          pure $ Regular $ last xs
        else 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
resType) (Segments -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
resType) NoMode
NoMode
          xs <- (ResRep -> FlattenM VName) -> [ResRep] -> 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 ResRep -> FlattenM VName
forall {f :: * -> *}. Applicative f => ResRep -> f VName
asRegular [ResRep]
branchesRep
          -- Create the blank space for the result
          resultSpace <- letExp "blank_res" =<< eBlank resultType
          -- Write back the values of each branch to the blank space
          result <- foldM (scatterRegular lvl) resultSpace $ zip iss xs
          result_t <- arrayShape <$> lookupType result
          result' <-
            letExp "match_res_reg" . BasicOp $
              Reshape result (reshapeAll result_t (segmentsShape segments <> arrayShape resType))
          pure $ Regular result'
  -- Irregular case
  | DistType Segments
_ Rank
_ (Array PrimType
pt Shape
_ NoMode
_) <- DistResult -> DistType
distResType DistResult
dist_res = do
      branchesIrregRep <- (ResRep -> FlattenM IrregularRep)
-> [ResRep] -> FlattenM [IrregularRep]
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 ResRep -> FlattenM IrregularRep
forall {f :: * -> *}. Applicative f => ResRep -> f IrregularRep
asIrregular [ResRep]
branchesRep
      let segsType = PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array (IntType -> PrimType
IntType IntType
Int64) (Segments -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) NoMode
NoMode
      -- Create a blank space for the 'segs'
      segsSpace <- letExp "blank_segs" =<< eBlank segsType
      -- Write back the segs of each branch to the blank space
      segs <- foldM (scatterRegular lvl) segsSpace $ zip iss (irregularS <$> branchesIrregRep)
      (_, offsets, num_data) <- exScanAndSum lvl segs
      let resultType = PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt (Segments -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
num_data]) NoMode
NoMode
      -- Create the blank space for the result
      resultSpace <- letExp "blank_res" =<< eBlank resultType
      -- Write back the values of each branch to the blank space
      elems <- foldM (scatterIrregular lvl offsets) resultSpace $ zip iss branchesIrregRep
      flags <- genFlags lvl num_data offsets
      pure $
        Irregular $
          IrregularRep
            { irregularS = segs,
              irregularF = flags,
              irregularO = offsets,
              irregularD = elems,
              irregularK = Dense
            }
  | Bool
otherwise = String -> FlattenM ResRep
forall a. HasCallStack => String -> a
error String
"mergeResult: non-array irregular result"
  where
    asRegular :: ResRep -> f VName
asRegular (Regular VName
v) = VName -> f VName
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
    asRegular ResRep
_ = String -> f VName
forall a. HasCallStack => String -> a
error String
"mergeResult: mismatched reps"

    asIrregular :: ResRep -> f IrregularRep
asIrregular (Irregular IrregularRep
irreg) = IrregularRep -> f IrregularRep
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IrregularRep
irreg
    asIrregular ResRep
_ = String -> f IrregularRep
forall a. HasCallStack => String -> a
error String
"mergeResult: mismatched reps"

-- | Flatten a single branch body of a variant 'Match', but guard its execution
-- on the branch actually being taken by some segment. When a branch receives no
-- segments (its partition is empty) we must not run its flattened code: it may
-- call lifted recursive functions, which would recurse forever on an empty
-- batch. An untaken branch's results are never read ('mergeResult' scatters
-- them back through the branch's empty index array), so we just yield blanks.
--
-- Like 'flattenUniformMatch', the branch is lifted to a 'Result' of flat rep
-- components and the reps recovered with 'distResultsToResReps'; here we
-- additionally wrap it in a @branch_size > 0@ 'Match'.
guardBranch ::
  FlattenOps ->
  SubExp ->
  DistEnv ->
  DistInputs ->
  DistStms ->
  [DistResult] ->
  Result ->
  FlattenM [ResRep]
guardBranch :: FlattenOps
-> SubExp
-> DistEnv
-> DistInputs
-> DistStms
-> [DistResult]
-> Result
-> FlattenM [ResRep]
guardBranch FlattenOps
ops SubExp
branch_size DistEnv
env DistInputs
inputs DistStms
dstms [DistResult]
res Result
result = do
  let branch_segments :: Segments
branch_segments = [SubExp
branch_size]
  (taken_body, taken_types) <-
    FlattenM (Result, [Type])
-> FlattenM (GBody (Rep FlattenM) SubExpRes, [Type])
forall (m :: * -> *) res a.
(MonadBuilder m, IsResult res) =>
m ([res], a) -> m (GBody (Rep m) res, a)
buildBody (FlattenM (Result, [Type])
 -> FlattenM (GBody (Rep FlattenM) SubExpRes, [Type]))
-> FlattenM (Result, [Type])
-> FlattenM (GBody (Rep FlattenM) SubExpRes, [Type])
forall a b. (a -> b) -> a -> b
$ do
      body_res <- FlattenOps
-> Segments
-> DistInputs
-> DistEnv
-> DistStms
-> [DistResult]
-> Result
-> FlattenM Result
liftBodyWithDistResults FlattenOps
ops Segments
branch_segments DistInputs
inputs DistEnv
env DistStms
dstms [DistResult]
res Result
result
      ts <- mapM (subExpType . resSubExp) body_res
      pure (body_res, ts)
  -- Blanks for the untaken branch have the same types as the taken branch.
  -- Sizes bound inside the branch (the length of irregular data) are not in
  -- scope here, so we zero them; the 'Match' then makes them existential.
  untaken_body <- buildBody_ $ do
    let blank TypeBase (ShapeBase a) m
t = TypeBase (ShapeBase a) m
t TypeBase (ShapeBase a) m -> Shape -> TypeBase Shape m
forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
`setArrayShape` Segments -> Shape
forall d. [d] -> ShapeBase d
Shape ((a -> SubExp) -> [a] -> Segments
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> a -> SubExp
forall a b. a -> b -> a
const (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0)) (TypeBase (ShapeBase a) m -> [a]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims TypeBase (ShapeBase a) m
t))
    subExpsRes <$> mapM (letSubExp "blank" <=< eBlank . blank) taken_types
  match_e <-
    eIf
      (eCmpOp (CmpSlt Int64) (eSubExp (intConst Int64 0)) (eSubExp branch_size))
      (pure taken_body)
      (pure untaken_body)
  match_res <- letTupExp "guarded_branch" match_e
  rets <- expExtType match_e
  pure $ distResultsToResReps res $ drop (S.size (shapeContext rets)) match_res

flattenVariantMatch ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  StmAux () ->
  [SubExp] ->
  [Case (Body SOACS)] ->
  Body SOACS ->
  MatchDec ExtType ->
  FlattenM DistEnv
flattenVariantMatch :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> Segments
-> [Case (Body SOACS)]
-> Body SOACS
-> MatchDec ExtType
-> FlattenM DistEnv
flattenVariantMatch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
_aux Segments
scrutinees [Case (Body SOACS)]
cases Body SOACS
defaultCase MatchDec ExtType
_rt = do
  let lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops
  w <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"w" (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
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ Segments -> [TPrimExp Int64 VName]
segmentDims Segments
segments
  -- We need to partition the indices of the scrutinees by which case they match.
  -- Lift the scrutinees.
  -- If it's a variable, we know it's a scalar and the lifted version will therefore be a regular array.
  lifted_scrutinees <- forM scrutinees $ \SubExp
scrut -> do
    SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl Segments
segments DistInputs
inps DistEnv
env (Segments -> Shape
segmentsShape Segments
segments) SubExp
scrut
  -- Cases for tagging values that match the same branch.
  -- The default case is the 0'th equvalence class.
  let equiv_cases =
        (Case (Body SOACS)
 -> Integer -> Case (FlattenM (GBody (Rep FlattenM) SubExpRes)))
-> [Case (Body SOACS)]
-> [Integer]
-> [Case (FlattenM (GBody (Rep FlattenM) SubExpRes))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          (\(Case [Maybe PrimValue]
pat Body SOACS
_) Integer
n -> [Maybe PrimValue]
-> FlattenM (GBody (Rep FlattenM) SubExpRes)
-> Case (FlattenM (GBody (Rep FlattenM) SubExpRes))
forall body. [Maybe PrimValue] -> body -> Case body
Case [Maybe PrimValue]
pat (FlattenM (GBody (Rep FlattenM) SubExpRes)
 -> Case (FlattenM (GBody (Rep FlattenM) SubExpRes)))
-> FlattenM (GBody (Rep FlattenM) SubExpRes)
-> Case (FlattenM (GBody (Rep FlattenM) SubExpRes))
forall a b. (a -> b) -> a -> b
$ [FlattenM (Exp (Rep FlattenM))]
-> FlattenM (GBody (Rep FlattenM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [SubExp -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp (SubExp -> FlattenM (Exp (Rep FlattenM)))
-> SubExp -> FlattenM (Exp (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
n])
          [Case (Body SOACS)]
cases
          [Integer
1 ..]
  let equiv_case_default = [FlattenM (Exp (Rep FlattenM))]
-> FlattenM (GBody (Rep FlattenM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [SubExp -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp (SubExp -> FlattenM (Exp (Rep FlattenM)))
-> SubExp -> FlattenM (Exp (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0]
  -- Match the scrutinees againts the branch cases
  equiv_classes <- letExp "equiv_classes" <=< segMap lvl (MkSolo w) $ \(MkSolo SubExp
i) -> do
    -- unflatten index
    let seg_is :: [TPrimExp Int64 VName]
seg_is = [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex (Segments -> [TPrimExp Int64 VName]
segmentDims Segments
segments) (SubExp -> TPrimExp Int64 VName
pe64 SubExp
i)
    scruts <- (VName -> FlattenM SubExp) -> [VName] -> FlattenM Segments
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"scruts" (Exp GPU -> FlattenM SubExp)
-> (VName -> FlattenM (Exp GPU)) -> VName -> FlattenM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< (VName -> [FlattenM (Exp GPU)] -> FlattenM (Exp GPU))
-> [FlattenM (Exp GPU)] -> VName -> FlattenM (Exp GPU)
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName
-> [FlattenM (Exp (Rep FlattenM))] -> FlattenM (Exp (Rep FlattenM))
VName -> [FlattenM (Exp GPU)] -> FlattenM (Exp GPU)
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex ((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]
seg_is)) [VName]
lifted_scrutinees
    cls <- letSubExp "cls" =<< eMatch scruts equiv_cases equiv_case_default
    pure [subExpRes cls]
  let num_cases = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [Case (Body SOACS)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Case (Body SOACS)]
cases Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  n_cases <- letExp "n_cases" <=< toExp $ intConst Int64 num_cases
  -- Parition the indices of the scrutinees by their equvalence class such
  -- that (the indices) of the scrutinees belonging to class 0 come first,
  -- then those belonging to class 1 and so on.
  (partition_sizes, partition_offs, partition_inds) <- doPartition lvl n_cases equiv_classes
  inds_t <- lookupType partition_inds
  -- Get the indices of each scrutinee by equivalence class
  branch_info <- forM [0 .. num_cases - 1] $ \Integer
i -> do
    num_data <-
      Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp (Name
"size" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> String -> Name
nameFromString (Integer -> String
forall a. Show a => a -> String
show Integer
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
partition_sizes [SubExp -> FlattenM (Exp (Rep FlattenM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp (SubExp -> FlattenM (Exp (Rep FlattenM)))
-> SubExp -> FlattenM (Exp (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
i]
    begin <-
      letSubExp ("idx_begin" <> nameFromString (show i))
        =<< eIndex partition_offs [toExp $ intConst Int64 i]
    inds <-
      letExp ("inds_branch" <> nameFromString (show i)) $
        BasicOp . Index partition_inds $
          fullSlice inds_t [DimSlice begin num_data (intConst Int64 1)]
    pure (num_data, inds)
  let (branch_sizes, inds) = unzip branch_info

  -- Distribute and lift the branch bodies.
  -- We put the default case at the start as it's the 0'th equivalence class
  -- and is therefore the first segment after the partition.
  let branch_bodies = Body SOACS
defaultCase Body SOACS -> [Body SOACS] -> [Body SOACS]
forall a. a -> [a] -> [a]
: (Case (Body SOACS) -> Body SOACS)
-> [Case (Body SOACS)] -> [Body SOACS]
forall a b. (a -> b) -> [a] -> [b]
map (\(Case [Maybe PrimValue]
_ Body SOACS
body) -> Body SOACS
body) [Case (Body SOACS)]
cases
  let branch_results = (Body SOACS -> Result) -> [Body SOACS] -> [Result]
forall a b. (a -> b) -> [a] -> [b]
map Body SOACS -> Result
forall rep res. GBody rep res -> [res]
bodyResult [Body SOACS]
branch_bodies
  -- Accumulator results are threaded from one branch to the next and cannot be
  -- blanked, so we only guard branch execution when no accumulators are
  -- involved. XXX: can we be sure this will never be a problem?
  let hasAcc = (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\DistResult
dr -> case DistResult -> DistType
distResType DistResult
dr of DistType Segments
_ Rank
_ Type
t -> Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t) [DistResult]
res
  -- acc inputs are handled differently, each branch use the result of the previous branch
  (branch_reps, _) <-
    foldM
      ( \([[ResRep]]
branch_reps_acc, Map VName ResRep
acc_reps) (SubExp
branch_size, VName
branch_inds, Body SOACS
body, Result
result) -> do
          let branch_segments :: Segments
branch_segments = [SubExp
branch_size]
          (inputs, env', dstms) <-
            FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> VName
-> Body SOACS
-> Map VName ResRep
-> FlattenM (DistInputs, DistEnv, DistStms)
distributeBranch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps VName
branch_inds Body SOACS
body Map VName ResRep
acc_reps
          reps <-
            if hasAcc
              then do
                env'' <- foldM (flattenDistStm ops branch_segments) env' dstms
                zipWithM (liftDistResultRep lvl branch_segments inputs env'') res result
              else guardBranch ops branch_size env' inputs dstms res result
          let acc_reps' = Map VName ResRep -> [ResRep] -> Map VName ResRep
forall {a}. Map VName a -> [a] -> Map VName a
replaceAccReps Map VName ResRep
acc_reps [ResRep]
reps
          pure (branch_reps_acc <> [reps], acc_reps')
      )
      ([], M.empty)
      (L.zip4 branch_sizes inds branch_bodies branch_results)
  -- Merging acc results is done by using the last branch result
  reps <- zipWithM (mergeResult lvl segments w inds) (L.transpose branch_reps) res
  insertRepsM (zip (map distResTag res) reps) env
  where
    findAccCert :: VName -> (VName, DistInput) -> Maybe VName
    findAccCert :: VName -> (VName, DistInput) -> Maybe VName
findAccCert VName
cert (VName, DistInput)
v_inp =
      let (VName
v, DistInput
inp) = (VName, DistInput)
v_inp
       in if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc (DistInput -> Type
distInputType DistInput
inp)
            then case DistInput -> Type
distInputType DistInput
inp of
              Acc VName
cert' Shape
_ [Type]
_ | VName
cert VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
cert' -> VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
              Type
_ -> Maybe VName
forall a. Maybe a
Nothing
            else Maybe VName
forall a. Maybe a
Nothing

    -- Idealy this should be a singleton
    findAccCerts :: VName -> [VName]
    findAccCerts :: VName -> [VName]
findAccCerts VName
cert = ((VName, DistInput) -> Maybe VName) -> DistInputs -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName -> (VName, DistInput) -> Maybe VName
findAccCert VName
cert) DistInputs
inps

    replaceAccRep :: Map VName a -> (DistResult, a) -> Map VName a
replaceAccRep Map VName a
acc_reps (DistResult
dist_res, a
rep) =
      let (DistType Segments
_ Rank
_ Type
t) = DistResult -> DistType
distResType DistResult
dist_res
       in if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
            then
              Map VName a
acc_reps
            else
              let (Acc VName
cert Shape
_ [Type]
_) = Type
t
                  accVars :: [VName]
accVars = VName -> [VName]
findAccCerts VName
cert
               in (Map VName a -> VName -> Map VName a)
-> Map VName a -> [VName] -> Map VName a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Map VName a
m VName
v -> VName -> a -> Map VName a -> Map VName a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v a
rep Map VName a
m) Map VName a
acc_reps [VName]
accVars
    replaceAccReps :: Map VName a -> [a] -> Map VName a
replaceAccReps Map VName a
acc_reps [a]
reps = (Map VName a -> (DistResult, a) -> Map VName a)
-> Map VName a -> [(DistResult, a)] -> Map VName a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map VName a -> (DistResult, a) -> Map VName a
forall {a}. Map VName a -> (DistResult, a) -> Map VName a
replaceAccRep Map VName a
acc_reps ([(DistResult, a)] -> Map VName a)
-> [(DistResult, a)] -> Map VName a
forall a b. (a -> b) -> a -> b
$ [DistResult] -> [a] -> [(DistResult, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [DistResult]
res [a]
reps

flattenUniformMatch ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  StmAux () ->
  [SubExp] ->
  [Case (Body SOACS)] ->
  Body SOACS ->
  MatchDec ExtType ->
  FlattenM DistEnv
flattenUniformMatch :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> Segments
-> [Case (Body SOACS)]
-> Body SOACS
-> MatchDec ExtType
-> FlattenM DistEnv
flattenUniformMatch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux Segments
scrutinees [Case (Body SOACS)]
cases Body SOACS
defaultCase MatchDec ExtType
rt = do
  scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  new_cases <- forM cases $ \(Case [Maybe PrimValue]
c Body SOACS
body) -> do
    let (DistInputs
case_body_inputs, DistStms
case_dstms) =
          FlattenOps
-> Scope GPU
-> Segments
-> DistInputs
-> Body SOACS
-> (DistInputs, DistStms)
forall rep.
FlattenOps
-> Scope rep
-> Segments
-> DistInputs
-> Body SOACS
-> (DistInputs, DistStms)
distributeBodyWith FlattenOps
ops Scope GPU
scope Segments
segments DistInputs
inps Body SOACS
body

    (GBody GPU SubExpRes -> Case (GBody GPU SubExpRes))
-> FlattenM (GBody GPU SubExpRes)
-> FlattenM (Case (GBody GPU SubExpRes))
forall a b. (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Maybe PrimValue]
-> GBody GPU SubExpRes -> Case (GBody GPU SubExpRes)
forall body. [Maybe PrimValue] -> body -> Case body
Case [Maybe PrimValue]
c) (FlattenM (GBody GPU SubExpRes)
 -> FlattenM (Case (GBody GPU SubExpRes)))
-> (FlattenM Result -> FlattenM (GBody GPU SubExpRes))
-> FlattenM Result
-> FlattenM (Case (GBody GPU SubExpRes))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlattenM Result -> FlattenM (GBody (Rep FlattenM) SubExpRes)
FlattenM Result -> FlattenM (GBody GPU SubExpRes)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (FlattenM Result -> FlattenM (Case (GBody GPU SubExpRes)))
-> FlattenM Result -> FlattenM (Case (GBody GPU SubExpRes))
forall a b. (a -> b) -> a -> b
$
      FlattenOps
-> Segments
-> DistInputs
-> DistEnv
-> DistStms
-> [DistResult]
-> Result
-> FlattenM Result
liftBodyWithDistResults FlattenOps
ops Segments
segments DistInputs
case_body_inputs DistEnv
env DistStms
case_dstms [DistResult]
res (Body SOACS -> Result
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
body)
  new_default_body <- do
    let (new_default_body_inputs, new_default_dstms) =
          distributeBodyWith ops scope segments inps defaultCase
    buildBody_ $
      liftBodyWithDistResults ops segments new_default_body_inputs env new_default_dstms res (bodyResult defaultCase)

  -- Maybe it is better to build MatchDec ourselves
  match_e <-
    eMatch'
      scrutinees
      [Case c (pure body) | Case c body <- new_cases]
      (pure new_default_body)
      (matchSort rt)

  match_res <-
    certifying (distCerts inps aux env) $
      letTupExp "match_res" match_e

  rets <- expExtType match_e
  -- get rid of the existential context
  let payload_res = Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
drop (Set Int -> Int
forall a. Set a -> Int
S.size ([ExtType] -> Set Int
forall o. [TypeBase (ShapeBase ExtSize) o] -> Set Int
shapeContext [ExtType]
rets)) [VName]
match_res
  let reps = [DistResult] -> [VName] -> [ResRep]
distResultsToResReps [DistResult]
res [VName]
payload_res
  insertRepsM (zip (map distResTag res) reps) env

-- | Flatten a 'Match'
flattenMatch ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  StmAux () ->
  [SubExp] ->
  [Case (Body SOACS)] ->
  Body SOACS ->
  MatchDec ExtType ->
  FlattenM DistEnv
flattenMatch :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> Segments
-> [Case (Body SOACS)]
-> Body SOACS
-> MatchDec ExtType
-> FlattenM DistEnv
flattenMatch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux Segments
scrutinees [Case (Body SOACS)]
cases Body SOACS
defaultCase MatchDec ExtType
rt =
  -- 'flattenUniformMatch' keeps the scrutinees in a plain GPU 'Match', which is
  -- only well-scoped when they are invariant to the nest. Whenever a scrutinee is
  -- variant we must partition the segments by branch, even if no branch contains
  -- parallelism (this happens e.g. for a variant conditional with an irregular
  -- result, which cannot be sequentialised into a scalar group).
  if (SubExp -> Bool) -> Segments -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> SubExp -> Bool
isVariant DistInputs
inps) Segments
scrutinees
    then FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> Segments
-> [Case (Body SOACS)]
-> Body SOACS
-> MatchDec ExtType
-> FlattenM DistEnv
flattenVariantMatch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux Segments
scrutinees [Case (Body SOACS)]
cases Body SOACS
defaultCase MatchDec ExtType
rt
    else FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> Segments
-> [Case (Body SOACS)]
-> Body SOACS
-> MatchDec ExtType
-> FlattenM DistEnv
flattenUniformMatch FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux Segments
scrutinees [Case (Body SOACS)]
cases Body SOACS
defaultCase MatchDec ExtType
rt