-- | Flattening of 'WithAcc'.
--
-- The basic idea is that in the nonuniform case, we change the 'WithAcc' to be
-- over the data array of the irregular representation. We then update all the
-- 'UpdateAcc' operations to compute flat indexes, via the usual metadata
-- arrays.
module Futhark.Pass.Flatten.WithAcc
  ( flattenWithAcc,
  )
where

import Control.Monad
import Data.Foldable
import Data.List qualified as L
import Data.Map qualified as M
import Futhark.IR.GPU
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Tools
import Futhark.Transform.ToGPU (soacsLambdaToGPU)
import Prelude hiding (div, rem)

indexIrreg ::
  (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
  Segments ->
  DistEnv ->
  IrregularRep ->
  SubExp ->
  Safety ->
  ShapeBase SubExp ->
  [SubExp] ->
  m SubExp
indexIrreg :: forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
Segments
-> DistEnv
-> IrregularRep
-> SubExp
-> Safety
-> Shape
-> Segments
-> m SubExp
indexIrreg Segments
_segments DistEnv
_env IrregularRep
rep SubExp
is Safety
safety Shape
shape Segments
js = do
  offset <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"uacc_segment_offset" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex (IrregularRep -> VName
irregularO IrregularRep
rep) [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is]
  flat <-
    letSubExp "flat_uacc_idx" <=< toExp $
      pe64 offset
        + flattenIndex
          (map pe64 (shapeDims shape))
          (map pe64 js)
  case safety of
    Safety
Unsafe -> SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
flat
    -- The original per-segment index 'js' may be out of bounds, which for a
    -- safe 'UpdateAcc' is a deliberate no-op. Adding the segment offset might
    -- turn such an out-of-bounds index into an in-bounds index into a
    -- neighbouring segment, so we must preserve the out-of-boundedness
    -- explicitly.
    Safety
Safe -> do
      inbounds <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"uacc_inbounds" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Shape -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
Shape -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eShapeInBounds Shape
shape ((SubExp -> m (Exp (Rep m))) -> Segments -> [m (Exp (Rep m))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp Segments
js)
      letSubExp "uacc_flat_idx"
        =<< eIf
          (eSubExp inbounds)
          (eBody [eSubExp flat])
          (eBody [eSubExp (intConst Int64 (-1))])

-- If just one input is nonuniform, we treat them all as nonuniform.
flattenWithAcc ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Pat Type ->
  StmAux () ->
  [WithAccInput SOACS] ->
  Lambda SOACS ->
  FlattenM DistEnv
flattenWithAcc :: FlattenOps
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat (TypeBase Shape NoMode)
-> StmAux ()
-> [WithAccInput SOACS]
-> Lambda SOACS
-> FlattenM DistEnv
flattenWithAcc FlattenOps
ops Segments
segments DistEnv
env DistInputs
inps [DistResult]
distres Pat (TypeBase Shape NoMode)
_withacc_pat StmAux ()
withacc_aux [WithAccInput SOACS]
withacc_inputs Lambda SOACS
acc_lam = do
  lam_params' <- [Param (TypeBase Shape NoMode)]
-> FlattenM [Param (TypeBase Shape NoMode)]
forall {m :: * -> *} {a}.
MonadFreshNames m =>
[Param a] -> m [Param a]
newAccLamParams ([Param (TypeBase Shape NoMode)]
 -> FlattenM [Param (TypeBase Shape NoMode)])
-> [Param (TypeBase Shape NoMode)]
-> FlattenM [Param (TypeBase Shape NoMode)]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
acc_lam

  iota_w <- genShapeIota (flattenSegLevel ops) $ segmentsShape segments

  iota_p <- newParam "iota_p" $ Prim int64
  -- Type in DistInputFree is parameter type
  let iota_w_t = PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  let iota_se = VName -> SubExp
Var (Param (TypeBase (ZonkAny 1) (ZonkAny 2)) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase (ZonkAny 1) (ZonkAny 2))
iota_p)

  -- Potentially change to distres option.
  let nonuniform =
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
          (WithAccInput SOACS -> Bool) -> [WithAccInput SOACS] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Shape
_, [VName]
arrs, Maybe (Lambda SOACS, Segments)
_) -> (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (DistEnv -> DistInputs -> VName -> Bool
isRegularInputArr DistEnv
env DistInputs
inps) [VName]
arrs) [WithAccInput SOACS]
withacc_inputs

  (withacc_inputs', trAccIndex, non_uniform_reps) <-
    if nonuniform
      then do
        (withacc_inputs', input_reps) <-
          mapAndUnzipM onNonuniformInput withacc_inputs
        let trAccIndex VName
c Safety
safety Segments
is = do
              ((shape, _, _), rep : _) <-
                VName
-> [(VName, (WithAccInput SOACS, [IrregularRep]))]
-> Maybe (WithAccInput SOACS, [IrregularRep])
forall a b. Eq a => a -> [(a, b)] -> Maybe b
L.lookup VName
c ([(VName, (WithAccInput SOACS, [IrregularRep]))]
 -> Maybe (WithAccInput SOACS, [IrregularRep]))
-> [(VName, (WithAccInput SOACS, [IrregularRep]))]
-> Maybe (WithAccInput SOACS, [IrregularRep])
forall a b. (a -> b) -> a -> b
$
                  [VName]
-> [(WithAccInput SOACS, [IrregularRep])]
-> [(VName, (WithAccInput SOACS, [IrregularRep]))]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (TypeBase Shape NoMode) -> VName)
-> [Param (TypeBase Shape NoMode)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
paramName [Param (TypeBase Shape NoMode)]
lam_params') ([(WithAccInput SOACS, [IrregularRep])]
 -> [(VName, (WithAccInput SOACS, [IrregularRep]))])
-> [(WithAccInput SOACS, [IrregularRep])]
-> [(VName, (WithAccInput SOACS, [IrregularRep]))]
forall a b. (a -> b) -> a -> b
$
                    [WithAccInput SOACS]
-> [[IrregularRep]] -> [(WithAccInput SOACS, [IrregularRep])]
forall a b. [a] -> [b] -> [(a, b)]
zip [WithAccInput SOACS]
withacc_inputs [[IrregularRep]]
input_reps
              Just $ L.singleton <$> indexIrreg segments env rep iota_se safety shape is
        pure (withacc_inputs', trAccIndex, concat input_reps)
      else do
        withacc_inputs' <- mapM onUniformInput withacc_inputs
        let trAccIndex VName
c p
_safety Segments
is = do
              _ <- VName
-> [(VName, WithAccInput SOACS)] -> Maybe (WithAccInput SOACS)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
L.lookup VName
c ([(VName, WithAccInput SOACS)] -> Maybe (WithAccInput SOACS))
-> [(VName, WithAccInput SOACS)] -> Maybe (WithAccInput SOACS)
forall a b. (a -> b) -> a -> b
$ [VName] -> [WithAccInput SOACS] -> [(VName, WithAccInput SOACS)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (TypeBase Shape NoMode) -> VName)
-> [Param (TypeBase Shape NoMode)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
paramName [Param (TypeBase Shape NoMode)]
lam_params') [WithAccInput SOACS]
withacc_inputs
              Just $ do
                iota_se_unflat <-
                  mapM (letSubExp "iota_idx" <=< toExp) $
                    unflattenIndex (segmentDims segments) (pe64 iota_se)
                pure $ iota_se_unflat ++ is
        pure (withacc_inputs', trAccIndex, [])
  let trAccShape VName
c = do
        (ispace, _, _) <- VName
-> [(VName, (Shape, [VName], Maybe (Lambda GPU, Segments)))]
-> Maybe (Shape, [VName], Maybe (Lambda GPU, Segments))
forall a b. Eq a => a -> [(a, b)] -> Maybe b
L.lookup VName
c ([(VName, (Shape, [VName], Maybe (Lambda GPU, Segments)))]
 -> Maybe (Shape, [VName], Maybe (Lambda GPU, Segments)))
-> [(VName, (Shape, [VName], Maybe (Lambda GPU, Segments)))]
-> Maybe (Shape, [VName], Maybe (Lambda GPU, Segments))
forall a b. (a -> b) -> a -> b
$ [VName]
-> [(Shape, [VName], Maybe (Lambda GPU, Segments))]
-> [(VName, (Shape, [VName], Maybe (Lambda GPU, Segments)))]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (TypeBase Shape NoMode) -> VName)
-> [Param (TypeBase Shape NoMode)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
paramName [Param (TypeBase Shape NoMode)]
lam_params') [(Shape, [VName], Maybe (Lambda GPU, Segments))]
withacc_inputs'
        pure ispace
      sf = (VName -> Maybe Shape
trAccShape, VName
-> Safety
-> Segments
-> Maybe (BuilderT SOACS (State VNameSource) Segments)
trAccIndex)

  acc_lam_body <-
    runBodyBuilder $
      localScope (scopeOfLParams lam_params') $
        bodyBind . lambdaBody =<< trLam sf acc_lam

  scope <- askScope
  let orig_acc_params = Int
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a. Int -> [a] -> [a]
drop Int
num_accs ([Param (TypeBase Shape NoMode)]
 -> [Param (TypeBase Shape NoMode)])
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
acc_lam
      lam_params_tr = (Param (TypeBase Shape NoMode) -> Param (TypeBase Shape NoMode))
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Param (TypeBase Shape NoMode) -> Param (TypeBase Shape NoMode)
forall o.
(VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Param (TypeBase Shape o) -> Param (TypeBase Shape o)
trParam (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf) [Param (TypeBase Shape NoMode)]
lam_params'
      acc_params_tr = Int
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a. Int -> [a] -> [a]
drop Int
num_accs [Param (TypeBase Shape NoMode)]
lam_params_tr
      interchanged_inps =
        (Param (TypeBase (ZonkAny 1) (ZonkAny 2)) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase (ZonkAny 1) (ZonkAny 2))
iota_p, VName -> TypeBase Shape NoMode -> DistInput
DistInputFree VName
iota_w TypeBase Shape NoMode
forall {shape} {o}. TypeBase shape o
iota_w_t)
          (VName, DistInput) -> DistInputs -> DistInputs
forall a. a -> [a] -> [a]
: [ (Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase Shape NoMode)
p, VName -> TypeBase Shape NoMode -> DistInput
DistInputFree (Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase Shape NoMode)
acc) (Param (TypeBase Shape NoMode) -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
paramType Param (TypeBase Shape NoMode)
acc))
            | -- This could potentially be wrong but since it's acc type it should be fine.
              (Param (TypeBase Shape NoMode)
p, Param (TypeBase Shape NoMode)
acc) <- [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
-> [(Param (TypeBase Shape NoMode), Param (TypeBase Shape NoMode))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape NoMode)]
orig_acc_params [Param (TypeBase Shape NoMode)]
acc_params_tr
            ]
          DistInputs -> DistInputs -> DistInputs
forall a. [a] -> [a] -> [a]
++ DistInputs
inps

  let (withacc_new_inputs, withacc_dstms) =
        distributeBodyWith ops scope segments interchanged_inps acc_lam_body

  withacc_lam' <- mkLambda (map (trParam sf) lam_params') $ do
    env' <- foldM (flattenDistStm ops segments) env withacc_dstms
    reps <-
      mapM
        (liftWithAccResult (flattenSegLevel ops) segments withacc_new_inputs env')
        (zip distres (bodyResult $ lambdaBody acc_lam))
    concat <$> mapM repToResults reps

  withacc_out_vs <-
    certifying (distCerts inps withacc_aux env) $
      letTupExp "withacc_flatten_out" (WithAcc withacc_inputs' withacc_lam')

  -- The accumulator results are handled differently in the nonuniform case,
  -- since we do not have metadata for them and since all of them are turned
  -- flat even when they might actually be regular. We can still turn the
  -- distres that are regular into Regulars here.
  let num_acc_results = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
arrs | (Shape
_, [VName]
arrs, Maybe (Lambda SOACS, Segments)
_) <- [WithAccInput SOACS]
withacc_inputs]
      (withacc_out_vs_wo, withacc_out_vs_no) = splitAt num_acc_results withacc_out_vs
      (distres_withacc, distres_normal) = splitAt num_acc_results distres

  let out_reps_normal = [DistResult] -> [VName] -> [ResRep]
resultToResRepsByDistResult [DistResult]
distres_normal [VName]
withacc_out_vs_no
  out_reps_withacc <-
    if nonuniform
      then mapM mkNonuniformWithAccRep (zip3 withacc_out_vs_wo non_uniform_reps distres_withacc)
      else pure $ map Regular withacc_out_vs_wo
  insertRepsM (zip (map distResTag $ distres_withacc ++ distres_normal) (out_reps_withacc ++ out_reps_normal)) env
  where
    newAccLamParams :: [Param a] -> m [Param a]
newAccLamParams [Param a]
ps = do
      let ([Param a]
cert_ps, [Param a]
acc_ps) = Int -> [Param a] -> ([Param a], [Param a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
num_accs [Param a]
ps
      -- Should not rename the certificates.
      acc_ps' <- [Param a] -> (Param a -> m (Param a)) -> m [Param a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param a]
acc_ps ((Param a -> m (Param a)) -> m [Param a])
-> (Param a -> m (Param a)) -> m [Param a]
forall a b. (a -> b) -> a -> b
$ \(Param Attrs
attrs VName
v a
t) ->
        Attrs -> VName -> a -> Param a
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
attrs (VName -> a -> Param a) -> m VName -> m (a -> Param a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
v m (a -> Param a) -> m a -> m (Param a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
t
      pure $ cert_ps <> acc_ps'

    num_accs :: Int
num_accs = [WithAccInput SOACS] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput SOACS]
withacc_inputs

    onOpWithIndexRank :: Int -> (Lambda SOACS, b) -> m (Lambda GPU, b)
onOpWithIndexRank Int
index_rank (Lambda SOACS
op_lam, b
nes) = do
      -- We need to add an additional index parameter because we are extending
      -- the index space of the accumulator. In the uniform case we have the
      -- full index space of the segments, while in the nonuniform case we only
      -- have one additional dimension.
      idx_ps <- Int
-> m (Param (TypeBase Shape NoMode))
-> m [Param (TypeBase Shape NoMode)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
index_rank (m (Param (TypeBase Shape NoMode))
 -> m [Param (TypeBase Shape NoMode)])
-> m (Param (TypeBase Shape NoMode))
-> m [Param (TypeBase Shape NoMode)]
forall a b. (a -> b) -> a -> b
$ Name -> TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"idx" (TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode)))
-> TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
      pure
        ( soacsLambdaToGPU $
            op_lam {lambdaParams = idx_ps <> lambdaParams op_lam},
          nes
        )

    -- Let's use liftSubExpRegular here
    onUniformInput :: (Shape, t VName, t (Lambda SOACS, b))
-> FlattenM (Shape, t VName, t (Lambda GPU, b))
onUniformInput (Shape
shape, t VName
arrs, t (Lambda SOACS, b)
op) =
      (Segments -> Shape
segmentsShape Segments
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
shape,,)
        (t VName
 -> t (Lambda GPU, b) -> (Shape, t VName, t (Lambda GPU, b)))
-> FlattenM (t VName)
-> FlattenM
     (t (Lambda GPU, b) -> (Shape, t VName, t (Lambda GPU, b)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> FlattenM VName) -> t VName -> FlattenM (t 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) -> t a -> m (t b)
mapM VName -> FlattenM VName
onArr t VName
arrs
        FlattenM (t (Lambda GPU, b) -> (Shape, t VName, t (Lambda GPU, b)))
-> FlattenM (t (Lambda GPU, b))
-> FlattenM (Shape, t VName, t (Lambda GPU, b))
forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((Lambda SOACS, b) -> FlattenM (Lambda GPU, b))
-> t (Lambda SOACS, b) -> FlattenM (t (Lambda GPU, b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse (Int -> (Lambda SOACS, b) -> FlattenM (Lambda GPU, b)
forall {m :: * -> *} {b}.
MonadFreshNames m =>
Int -> (Lambda SOACS, b) -> m (Lambda GPU, b)
onOpWithIndexRank (Segments -> Int
segmentsRank Segments
segments)) t (Lambda SOACS, b)
op
      where
        onArr :: VName -> FlattenM VName
onArr VName
arr = do
          arr_t <- DistInputs -> VName -> FlattenM (TypeBase Shape NoMode)
lookupInputType DistInputs
inps VName
arr
          let arr_shape = TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase Shape NoMode
arr_t
              expected_shape = Segments -> Shape
segmentsShape Segments
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
arr_shape
          liftSubExpRegular (flattenSegLevel ops) segments inps env expected_shape (Var arr)

    onNonuniformOp :: Int -> (Lambda SOACS, b) -> m (Lambda GPU, b)
onNonuniformOp Int
rank (Lambda SOACS
op, b
nes) = do
      let ([LParam SOACS]
old_idx_ps, [LParam SOACS]
value_ps) = Int -> [LParam SOACS] -> ([LParam SOACS], [LParam SOACS])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
rank ([LParam SOACS] -> ([LParam SOACS], [LParam SOACS]))
-> [LParam SOACS] -> ([LParam SOACS], [LParam SOACS])
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
op
          old_indices_used :: Bool
old_indices_used =
            (LParam SOACS -> Bool) -> [LParam SOACS] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\LParam SOACS
p -> LParam SOACS -> VName
forall dec. Param dec -> VName
paramName LParam SOACS
p VName -> Names -> Bool
`nameIn` GBody SOACS SubExpRes -> Names
forall a. FreeIn a => a -> Names
freeIn (Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
op)) [LParam SOACS]
old_idx_ps
      -- XXX: It is possible to change the lambda body to restore the indices based on flat_idx_p and handle this.
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
old_indices_used (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
        String -> m ()
forall a. HasCallStack => String -> a
error String
"flattenWithAcc: accumulator operator uses nonuniform indices"
      flat_idx_p <- Name -> TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"flat_idx" (TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode)))
-> TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
      pure
        ( soacsLambdaToGPU $
            op
              { lambdaParams = flat_idx_p : value_ps
              },
          nes
        )
    onNonuniformInput :: (a, [VName], t (Lambda SOACS, b))
-> FlattenM ((Shape, [VName], t (Lambda GPU, b)), [IrregularRep])
onNonuniformInput (a
shape, [VName]
arrs, t (Lambda SOACS, b)
op) = do
      reps <- (VName -> FlattenM IrregularRep)
-> [VName] -> 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 (SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> VName
-> FlattenM IrregularRep
getIrregRep (FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops) Segments
segments DistEnv
env DistInputs
inps) [VName]
arrs
      -- We need to ensure that the irregular arrays are dense
      reps_dense <- mapM (ensureDenseIrregular (flattenSegLevel ops) "withacc_input") reps
      let arrs' = (IrregularRep -> VName) -> [IrregularRep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map IrregularRep -> VName
irregularD [IrregularRep]
reps_dense
      w <- fmap (arraySize 0) . lookupType $ head arrs'
      -- We need to reshape to make sure all of the inputs have the same shape
      arrs'' <- forM arrs' $ \VName
v -> do
        v_t <- VName -> FlattenM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
        letExp (baseName v <> "_withacc_input_reshaped") . BasicOp $
          Reshape v $
            reshapeAll (arrayShape v_t) (Shape [w])
      (,reps_dense) . (Shape [w],arrs'',) <$> traverse (onNonuniformOp (shapeRank shape)) op

    -- The irregular kind is not carried through the results of the
    -- WithAcc, and 'mkNormalResReps' reconstructs the rep as Dense, so
    -- any irregular rep must actually be made dense before it crosses
    -- the WithAcc boundary.
    ensureDenseRep :: SegLevel -> ResRep -> FlattenM ResRep
ensureDenseRep SegLevel
lvl (Irregular IrregularRep
irreg) =
      IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep)
-> FlattenM IrregularRep -> FlattenM ResRep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep
ensureDenseIrregular SegLevel
lvl Name
"withacc_result" IrregularRep
irreg
    ensureDenseRep SegLevel
_ ResRep
rep = ResRep -> FlattenM ResRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ResRep
rep

    liftWithAccResult :: SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> (DistResult, SubExpRes)
-> FlattenM ResRep
liftWithAccResult SegLevel
lvl Segments
segs DistInputs
inputs DistEnv
env' (DistResult
dist_res, SubExpRes
res) =
      case SubExpRes -> SubExp
resSubExp SubExpRes
res of
        Var VName
v -> do
          let (Just (TypeBase Shape NoMode
t, ResRep
rep)) = VName
-> Map VName (TypeBase Shape NoMode, ResRep)
-> Maybe (TypeBase Shape NoMode, ResRep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName (TypeBase Shape NoMode, ResRep)
 -> Maybe (TypeBase Shape NoMode, ResRep))
-> Map VName (TypeBase Shape NoMode, ResRep)
-> Maybe (TypeBase Shape NoMode, ResRep)
forall a b. (a -> b) -> a -> b
$ DistInputs -> DistEnv -> Map VName (TypeBase Shape NoMode, ResRep)
inputReps DistInputs
inputs DistEnv
env'
          if TypeBase Shape NoMode -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc TypeBase Shape NoMode
t
            then
              ResRep -> FlattenM ResRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ResRep
rep
            else
              SegLevel -> ResRep -> FlattenM ResRep
ensureDenseRep SegLevel
lvl (ResRep -> FlattenM ResRep) -> FlattenM ResRep -> FlattenM ResRep
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> DistResult
-> SubExpRes
-> FlattenM ResRep
liftDistResultRep SegLevel
lvl Segments
segs DistInputs
inputs DistEnv
env' DistResult
dist_res SubExpRes
res
        Constant PrimValue
_ -> SegLevel -> ResRep -> FlattenM ResRep
ensureDenseRep SegLevel
lvl (ResRep -> FlattenM ResRep) -> FlattenM ResRep -> FlattenM ResRep
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SegLevel
-> Segments
-> DistInputs
-> DistEnv
-> DistResult
-> SubExpRes
-> FlattenM ResRep
liftDistResultRep SegLevel
lvl Segments
segs DistInputs
inputs DistEnv
env' DistResult
dist_res SubExpRes
res

    repToResults :: ResRep -> f [SubExpRes]
repToResults (Regular VName
v) =
      [SubExpRes] -> f [SubExpRes]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Certs -> SubExp -> SubExpRes
SubExpRes Certs
forall a. Monoid a => a
mempty (SubExp -> SubExpRes) -> SubExp -> SubExpRes
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v]
    repToResults (Irregular IrregularRep
irreg) =
      (VName -> SubExpRes) -> [VName] -> [SubExpRes]
forall a b. (a -> b) -> [a] -> [b]
map (Certs -> SubExp -> SubExpRes
SubExpRes Certs
forall a. Monoid a => a
mempty (SubExp -> SubExpRes) -> (VName -> SubExp) -> VName -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) ([VName] -> [SubExpRes]) -> f [VName] -> f [SubExpRes]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IrregularRep -> f [VName]
forall {m :: * -> *}. MonadBuilder m => IrregularRep -> m [VName]
irregResults IrregularRep
irreg

    irregResults :: IrregularRep -> m [VName]
irregResults
      ( IrregularRep
          { irregularS :: IrregularRep -> VName
irregularS = VName
segs,
            irregularF :: IrregularRep -> VName
irregularF = VName
flags,
            irregularO :: IrregularRep -> VName
irregularO = VName
offsets,
            irregularD :: IrregularRep -> VName
irregularD = VName
elems
          }
        ) = do
        flags_t <- VName -> m (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
flags
        t <- lookupType elems
        num_data <- letExp "num_data" =<< toExp (product $ map pe64 $ arrayDims t)
        let shape = Segments -> Shape
forall d. [d] -> ShapeBase d
Shape [VName -> SubExp
Var VName
num_data]
        flags' <- letExp "flags" $ BasicOp $ Reshape flags $ reshapeAll (arrayShape flags_t) shape
        elems' <- letExp "elems" $ BasicOp $ Reshape elems $ reshapeAll (arrayShape t) shape
        pure [num_data, segs, flags', offsets, elems']

    mkNonuniformWithAccRep :: (VName, IrregularRep, DistResult) -> m ResRep
mkNonuniformWithAccRep (VName
v, IrregularRep
rep, DistResult
dist_res)
      | DistResult -> Bool
isRegularDistResult DistResult
dist_res = do
          let DistType Segments
_ Rank
_ TypeBase Shape NoMode
t = DistResult -> DistType
distResType DistResult
dist_res
              expectedShape :: Shape
expectedShape = Segments -> Shape
segmentsShape Segments
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase Shape NoMode
t
          v_t <- VName -> m (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
          v_reshaped <-
            letExp "actual_regular_with_acc_res" . BasicOp $
              Reshape v (reshapeAll (arrayShape v_t) expectedShape)
          pure $ Regular v_reshaped
      | Bool
otherwise =
          ResRep -> m ResRep
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResRep -> m ResRep) -> ResRep -> m ResRep
forall a b. (a -> b) -> a -> b
$ IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep) -> IrregularRep -> ResRep
forall a b. (a -> b) -> a -> b
$ IrregularRep
rep {irregularD = v}

    trType ::
      (VName -> Maybe Shape, VName -> Safety -> [SubExp] -> Maybe (Builder SOACS [SubExp])) ->
      TypeBase shape o ->
      TypeBase shape o
    trType :: forall shape o.
(VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> TypeBase shape o -> TypeBase shape o
trType (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf (Acc VName
acc Shape
_ [TypeBase Shape NoMode]
ts)
      | Just Shape
shape <- (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> VName -> Maybe Shape
forall a b. (a, b) -> a
fst (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf VName
acc =
          VName -> Shape -> [TypeBase Shape NoMode] -> TypeBase shape o
forall shape o.
VName -> Shape -> [TypeBase Shape NoMode] -> TypeBase shape o
Acc VName
acc Shape
shape [TypeBase Shape NoMode]
ts
    trType (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
_ TypeBase shape o
t = TypeBase shape o
t

    trParam ::
      (VName -> Maybe Shape, VName -> Safety -> [SubExp] -> Maybe (Builder SOACS [SubExp])) ->
      Param (TypeBase Shape o) ->
      Param (TypeBase Shape o)
    trParam :: forall o.
(VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Param (TypeBase Shape o) -> Param (TypeBase Shape o)
trParam (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf = (TypeBase Shape o -> TypeBase Shape o)
-> Param (TypeBase Shape o) -> Param (TypeBase Shape o)
forall a b. (a -> b) -> Param a -> Param b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase Shape o -> TypeBase Shape o)
 -> Param (TypeBase Shape o) -> Param (TypeBase Shape o))
-> (TypeBase Shape o -> TypeBase Shape o)
-> Param (TypeBase Shape o)
-> Param (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> TypeBase Shape o -> TypeBase Shape o
forall shape o.
(VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> TypeBase shape o -> TypeBase shape o
trType (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf

    trBody :: (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> GBody SOACS SubExpRes
-> BuilderT SOACS (State VNameSource) (GBody SOACS SubExpRes)
trBody (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf (Body BodyDec SOACS
dec Stms SOACS
stms [SubExpRes]
res) =
      BodyDec SOACS -> Stms SOACS -> [SubExpRes] -> GBody SOACS SubExpRes
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body BodyDec SOACS
dec (Stms SOACS -> [SubExpRes] -> GBody SOACS SubExpRes)
-> BuilderT SOACS (State VNameSource) (Stms SOACS)
-> BuilderT
     SOACS (State VNameSource) ([SubExpRes] -> GBody SOACS SubExpRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BuilderT SOACS (State VNameSource) ()
-> BuilderT
     SOACS
     (State VNameSource)
     (Stms (Rep (BuilderT SOACS (State VNameSource))))
forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
collectStms_ ((Stm SOACS -> BuilderT SOACS (State VNameSource) ())
-> Stms SOACS -> BuilderT SOACS (State VNameSource) ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Stm SOACS -> BuilderT SOACS (State VNameSource) ()
onStm Stms SOACS
stms) BuilderT
  SOACS (State VNameSource) ([SubExpRes] -> GBody SOACS SubExpRes)
-> Builder SOACS [SubExpRes]
-> BuilderT SOACS (State VNameSource) (GBody SOACS SubExpRes)
forall a b.
BuilderT SOACS (State VNameSource) (a -> b)
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [SubExpRes] -> Builder SOACS [SubExpRes]
forall a. a -> BuilderT SOACS (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExpRes]
res
      where
        onStm :: Stm SOACS -> BuilderT SOACS (State VNameSource) ()
onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) =
          Stm (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
Stm SOACS -> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm SOACS -> BuilderT SOACS (State VNameSource) ())
-> (Exp SOACS -> Stm SOACS)
-> Exp SOACS
-> BuilderT SOACS (State VNameSource) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ((TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> Pat (TypeBase Shape NoMode) -> Pat (TypeBase Shape NoMode)
forall a b. (a -> b) -> Pat a -> Pat b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall shape o.
(VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> TypeBase shape o -> TypeBase shape o
trType (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf) Pat (TypeBase Shape NoMode)
Pat (LetDec SOACS)
pat) StmAux (ExpDec SOACS)
aux (Exp SOACS -> BuilderT SOACS (State VNameSource) ())
-> BuilderT SOACS (State VNameSource) (Exp SOACS)
-> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Pat (TypeBase Shape NoMode)
-> Exp SOACS
-> BuilderT SOACS (State VNameSource) (Exp SOACS)
trExp (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf Pat (TypeBase Shape NoMode)
Pat (LetDec SOACS)
pat Exp SOACS
e

    trLam :: (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Lambda SOACS
-> BuilderT SOACS (State VNameSource) (Lambda SOACS)
trLam (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf (Lambda [LParam SOACS]
params [TypeBase Shape NoMode]
ret GBody SOACS SubExpRes
body) =
      [LParam SOACS]
-> [TypeBase Shape NoMode] -> GBody SOACS SubExpRes -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda ((Param (TypeBase Shape NoMode) -> Param (TypeBase Shape NoMode))
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Param (TypeBase Shape NoMode) -> Param (TypeBase Shape NoMode)
forall o.
(VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Param (TypeBase Shape o) -> Param (TypeBase Shape o)
trParam (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf) [Param (TypeBase Shape NoMode)]
[LParam SOACS]
params) ((TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall shape o.
(VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> TypeBase shape o -> TypeBase shape o
trType (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf) [TypeBase Shape NoMode]
ret) (GBody SOACS SubExpRes -> Lambda SOACS)
-> BuilderT SOACS (State VNameSource) (GBody SOACS SubExpRes)
-> BuilderT SOACS (State VNameSource) (Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> GBody SOACS SubExpRes
-> BuilderT SOACS (State VNameSource) (GBody SOACS SubExpRes)
trBody (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf GBody SOACS SubExpRes
body

    trSOAC :: (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Op SOACS -> BuilderT SOACS (State VNameSource) (Op SOACS)
trSOAC (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf = SOACMapper SOACS SOACS (BuilderT SOACS (State VNameSource))
-> SOAC SOACS -> BuilderT SOACS (State VNameSource) (SOAC SOACS)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper SOACS SOACS (BuilderT SOACS (State VNameSource))
mapper
      where
        mapper :: SOACMapper SOACS SOACS (BuilderT SOACS (State VNameSource))
mapper =
          SOACMapper SOACS SOACS (BuilderT SOACS (State VNameSource))
forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper {mapOnSOACLambda = trLam sf}

    trExp :: (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Pat (TypeBase Shape NoMode)
-> Exp SOACS
-> BuilderT SOACS (State VNameSource) (Exp SOACS)
trExp (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf Pat (TypeBase Shape NoMode)
_ (WithAcc [WithAccInput SOACS]
acc_inputs Lambda SOACS
lam) =
      [WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [WithAccInput SOACS]
acc_inputs (Lambda SOACS -> Exp SOACS)
-> BuilderT SOACS (State VNameSource) (Lambda SOACS)
-> BuilderT SOACS (State VNameSource) (Exp SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> Lambda SOACS
-> BuilderT SOACS (State VNameSource) (Lambda SOACS)
trLam (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf Lambda SOACS
lam
    trExp (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf (Pat [PatElem VName
_ TypeBase Shape NoMode
acc_t]) (BasicOp (UpdateAcc Safety
safety VName
acc Segments
is Segments
ses)) = do
      case TypeBase Shape NoMode
acc_t of
        Acc VName
cert Shape
_ [TypeBase Shape NoMode]
_
          | Just BuilderT SOACS (State VNameSource) Segments
mk <- (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> VName
-> Safety
-> Segments
-> Maybe (BuilderT SOACS (State VNameSource) Segments)
forall a b. (a, b) -> b
snd (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf VName
cert Safety
safety Segments
is -> do
              is' <- BuilderT SOACS (State VNameSource) Segments
mk
              pure $ BasicOp $ UpdateAcc safety acc is' ses
        TypeBase Shape NoMode
_ ->
          Exp SOACS -> BuilderT SOACS (State VNameSource) (Exp SOACS)
forall a. a -> BuilderT SOACS (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp SOACS -> BuilderT SOACS (State VNameSource) (Exp SOACS))
-> Exp SOACS -> BuilderT SOACS (State VNameSource) (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ Safety -> VName -> Segments -> Segments -> BasicOp
UpdateAcc Safety
safety VName
acc Segments
is Segments
ses
    trExp (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf Pat (TypeBase Shape NoMode)
_ Exp SOACS
e = Mapper SOACS SOACS (BuilderT SOACS (State VNameSource))
-> Exp SOACS -> BuilderT SOACS (State VNameSource) (Exp SOACS)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS (BuilderT SOACS (State VNameSource))
mapper Exp SOACS
e
      where
        mapper :: Mapper SOACS SOACS (BuilderT SOACS (State VNameSource))
mapper =
          Mapper (ZonkAny 0) (ZonkAny 0) (BuilderT SOACS (State VNameSource))
forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper
            { mapOnBody = \Scope SOACS
_ -> (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
-> GBody SOACS SubExpRes
-> BuilderT SOACS (State VNameSource) (GBody SOACS SubExpRes)
trBody (VName -> Maybe Shape,
 VName
 -> Safety
 -> Segments
 -> Maybe (BuilderT SOACS (State VNameSource) Segments))
sf,
              mapOnRetType = pure . trType sf,
              mapOnBranchType = pure . trType sf,
              mapOnFParam = pure . trParam sf,
              mapOnLParam = pure . trParam sf,
              mapOnOp = trSOAC sf
            }