module Futhark.Pass.Flatten.Loop
  ( flattenLoop,
  )
where

import Control.Monad
import Control.Monad.Reader (runReaderT)
import Data.Containers.ListUtils (nubOrd)
import Data.Foldable
import Data.List qualified as L
import Data.Map qualified as M
import Data.Set qualified as S
import Data.Tuple.Solo
import Futhark.IR.GPU
import Futhark.IR.SOACS
import Futhark.IR.SOACS.Simplify (simplifyStms)
import Futhark.MonadFreshNames
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Tools
import Prelude hiding (div, quot, rem)

-- | Is this dimension variant to the loop or the outer map context -
-- either because it is itself a loop parameter, or because it is
-- variant in the outer map-nest?
variantDim :: DistInputs -> S.Set VName -> SubExp -> Bool
variantDim :: DistInputs -> Set VName -> SubExp -> Bool
variantDim DistInputs
_ Set VName
_ Constant {} = Bool
False
variantDim DistInputs
inps Set VName
loopParamNames (Var VName
v) =
  VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
loopParamNames Bool -> Bool -> Bool
|| DistInputs -> SubExp -> Bool
isVariant DistInputs
inps (VName -> SubExp
Var VName
v)

-- Check whether a loop parameter array needs irregular representation.
-- we need the irregular representation when any of its dimensions are either:
-- a loop parameter name or variant in the outer map context
needsIrregular :: DistInputs -> S.Set VName -> DeclType -> Bool
needsIrregular :: DistInputs -> Set VName -> DeclType -> Bool
needsIrregular DistInputs
inps Set VName
loopParamNames DeclType
t =
  case DeclType
t of
    Array {} -> (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> Set VName -> SubExp -> Bool
variantDim DistInputs
inps Set VName
loopParamNames) (DeclType -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims DeclType
t)
    DeclType
_ -> Bool
False

-- Lift a loop parameter and its initial value together.
-- If the parameter is an array whose dimensions are all invariant,
-- we lift it to a regular array. Otherwise we fall back to irregular.
liftLoopParam ::
  SegLevel ->
  Segments ->
  SubExp ->
  DistInputs ->
  DistEnv ->
  S.Set VName ->
  (FParam SOACS, SubExp) ->
  FlattenM ([FParam GPU], ResRep, [SubExp])
liftLoopParam :: SegLevel
-> [SubExp]
-> SubExp
-> DistInputs
-> DistEnv
-> Set VName
-> (FParam SOACS, SubExp)
-> FlattenM ([FParam GPU], ResRep, [SubExp])
liftLoopParam SegLevel
lvl [SubExp]
segments SubExp
num_segments DistInputs
inps DistEnv
env Set VName
loopParamNames (FParam SOACS
fparam, SubExp
initSE) = do
  let t :: DeclType
t = Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
FParam SOACS
fparam
  case DeclType
t of
    Prim PrimType
pt -> do
      param <-
        Name -> DeclType -> FlattenM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam
          (VName -> Name
baseName (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
FParam SOACS
fparam) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_lifted")
          (TypeBase Shape (ZonkAny 2) -> Shape -> Diet -> DeclType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase Shape (ZonkAny 2)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) ([SubExp] -> Shape
segmentsShape [SubExp]
segments) Diet
Observe)
      initV <- liftSubExpRegular lvl segments inps env (segmentsShape segments) initSE
      pure ([param], Regular $ paramName param, [Var initV])
    Array PrimType
pt Shape
_ Diet
o
      | DistInputs -> Set VName -> DeclType -> Bool
needsIrregular DistInputs
inps Set VName
loopParamNames DeclType
t -> do
          (params, rep) <- SubExp -> FParam SOACS -> FlattenM ([FParam GPU], ResRep)
forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> FParam SOACS -> m ([FParam GPU], ResRep)
liftParam SubExp
num_segments FParam SOACS
fparam
          (_, initRep) <- liftSubExp lvl segments inps env initSE
          irreg <- case initRep of
            -- This will not happen.
            Regular VName
v -> SegLevel -> [SubExp] -> VName -> FlattenM IrregularRep
mkIrregFromReg SegLevel
lvl [SubExp]
segments VName
v
            Irregular IrregularRep
irreg -> IrregularRep -> FlattenM IrregularRep
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IrregularRep
irreg
          initVals <- irregularRepToFlatArrs num_segments irreg
          pure (params, rep, map Var initVals)
      | Bool
otherwise -> do
          -- Regular case: all dims are invariant, just add w as outermost dim
          let pShape :: Shape
pShape = [SubExp] -> Shape
segmentsShape [SubExp]
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> DeclType -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape DeclType
t
          p <-
            Name -> DeclType -> FlattenM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam
              (VName -> Name
baseName (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
FParam SOACS
fparam) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_lifted")
              (TypeBase Shape (ZonkAny 3) -> Shape -> Diet -> DeclType
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (PrimType -> TypeBase Shape (ZonkAny 3)
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt) Shape
pShape Diet
o)
          initV <- liftSubExpRegular lvl segments inps env pShape initSE
          -- If the parameter is consumed, we must not consume the
          -- representation array (it may be used by other versions in
          -- multi-versioned code), so insert a copy. The simplifier hopefully
          -- removes it again when consuming the representation directly is
          -- safe.
          initV' <-
            if o == Consume
              then letExp (baseName (paramName fparam) <> "_inter_copy") =<< eCopy (eVar initV)
              else pure initV
          pure ([p], Regular $ paramName p, [Var initV'])
    Acc {} -> do
      initV <- SegLevel
-> [SubExp]
-> DistInputs
-> DistEnv
-> Shape
-> SubExp
-> FlattenM VName
liftSubExpRegular SegLevel
lvl [SubExp]
segments DistInputs
inps DistEnv
env Shape
forall a. Monoid a => a
mempty SubExp
initSE
      let Param attrs v acc_t = fparam
      param <- Param attrs <$> newName v <*> pure acc_t
      pure ([param], Regular $ paramName param, [Var initV])
    Mem {} ->
      String -> FlattenM ([Param DeclType], ResRep, [SubExp])
forall a. HasCallStack => String -> a
error String
"liftLoopParam: Mem"

-- | Construct the body of an interchanged uniform loop: a single
-- Screma mapping the original loop body over the lifted loop
-- parameters (and any other inputs used by the body), transformed as
-- if it were a top-level statement - in particular, it is subject to
-- multi-versioning. The lambda parameters reuse the original names,
-- so the body can be used unchanged. Only usable when all involved
-- values are regular.
interchangedLoopBody ::
  FlattenOps ->
  SubExp ->
  Segments ->
  DistEnv ->
  [(FParam SOACS, FParam GPU)] ->
  DistInputs ->
  StmAux () ->
  Body SOACS ->
  FlattenM (Body GPU)
interchangedLoopBody :: FlattenOps
-> SubExp
-> [SubExp]
-> DistEnv
-> [(FParam SOACS, FParam GPU)]
-> DistInputs
-> StmAux ()
-> Body SOACS
-> FlattenM (Body GPU)
interchangedLoopBody FlattenOps
ops SubExp
num_segments [SubExp]
segments DistEnv
env [(FParam SOACS, FParam GPU)]
params DistInputs
free_inps StmAux ()
aux Body SOACS
body = FlattenM [SubExpRes] -> FlattenM (GBody (Rep FlattenM) SubExpRes)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (FlattenM [SubExpRes] -> FlattenM (GBody (Rep FlattenM) SubExpRes))
-> FlattenM [SubExpRes]
-> FlattenM (GBody (Rep FlattenM) SubExpRes)
forall a b. (a -> b) -> a -> b
$ do
  let flatInput :: VName -> VName -> TypeBase Shape o -> m VName
flatInput VName
name VName
arr TypeBase Shape o
t = do
        arr_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
        letExp (baseName name <> "_flat") . BasicOp . Reshape arr $
          reshapeAll (arrayShape arr_t) (Shape [num_segments] <> arrayShape t)
      inputArr :: DistInput -> VName
inputArr (DistInputFree VName
arr Type
_) = VName
arr
      inputArr (DistInput ResTag
rt Type
_) = case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
        Regular VName
arr -> VName
arr
        Irregular {} -> String -> VName
forall a. HasCallStack => String -> a
error String
"interchangedLoopBody: irregular input"
  param_arrs <- [(Param DeclType, Param DeclType)]
-> ((Param DeclType, Param DeclType) -> FlattenM VName)
-> FlattenM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Param DeclType, Param DeclType)]
[(FParam SOACS, FParam GPU)]
params (((Param DeclType, Param DeclType) -> FlattenM VName)
 -> FlattenM [VName])
-> ((Param DeclType, Param DeclType) -> FlattenM VName)
-> FlattenM [VName]
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, Param DeclType
lifted_p) ->
    VName -> VName -> Type -> FlattenM VName
forall {m :: * -> *} {o}.
MonadBuilder m =>
VName -> VName -> TypeBase Shape o -> m VName
flatInput (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p) (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
lifted_p) (DeclType -> Type
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl (Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
p))
  free_arrs <- forM free_inps $ \(VName
v, DistInput
inp) ->
    VName -> VName -> Type -> FlattenM VName
forall {m :: * -> *} {o}.
MonadBuilder m =>
VName -> VName -> TypeBase Shape o -> m VName
flatInput VName
v (DistInput -> VName
inputArr DistInput
inp) (DistInput -> Type
distInputType DistInput
inp)

  let lam_params =
        [Attrs -> VName -> Type -> Param Type
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p) (DeclType -> Type
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl (Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
p)) | (Param DeclType
p, Param DeclType
_) <- [(Param DeclType, Param DeclType)]
[(FParam SOACS, FParam GPU)]
params]
          [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Attrs -> VName -> Type -> Param Type
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
v (DistInput -> Type
distInputType DistInput
inp) | (VName
v, DistInput
inp) <- DistInputs
free_inps]
      row_ts = [DeclType -> Type
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl (Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
p) | (Param DeclType
p, Param DeclType
_) <- [(Param DeclType, Param DeclType)]
[(FParam SOACS, FParam GPU)]
params]
      lam = [LParam SOACS] -> [Type] -> Body SOACS -> GLambda SOACS Type
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
lam_params [Type]
row_ts Body SOACS
body
  pes <- forM (zip params row_ts) $ \((Param DeclType
p, Param DeclType
_), Type
t) ->
    VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem
      (VName -> Type -> PatElem Type)
-> FlattenM VName -> FlattenM (Type -> PatElem Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlattenM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)
      FlattenM (Type -> PatElem Type)
-> FlattenM Type -> FlattenM (PatElem Type)
forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> FlattenM Type
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
t Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
`arrayOfRow` SubExp
num_segments)
  form <- mapSOAC lam
  let map_stm :: Stm SOACS
      map_stm =
        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]
pes) (StmAux ()
aux {stmAuxCerts = mempty}) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
          Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
num_segments ([VName]
param_arrs [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
free_arrs) ScremaForm SOACS
form

  -- Simplify before transforming. Apart from generally producing
  -- better code, this hoists statements that are invariant to the
  -- mapped values out of the Screma, and in particular any sizes
  -- they compute must be in scope when the Screma is versioned
  -- (e.g. for deciding intrablock feasibility).
  scope <- castScope <$> askScope
  map_stms <- runReaderT (simplifyStms (oneStm map_stm)) (scope :: Scope SOACS)
  mapM_ (flattenTopLevelStm ops) map_stms
  fmap (map (SubExpRes mempty . Var)) . forM (zip pes params) $ \(PatElem Type
pe, (Param DeclType
p, Param DeclType
_)) -> do
    pe_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType (PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe)
    let seg_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 (DeclType -> Type
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl (Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
p))
    letExp (baseName (paramName p) <> "_unflat") . BasicOp . Reshape (patElemName pe) $
      reshapeAll (arrayShape pe_t) seg_shape

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

-- | Distribute the loop body statement by statement and lift the
-- distributed statements, producing the statements and result of the
-- body of the lifted loop. The provided scope is that of the lifted
-- loop parameters (and any loop index); it is brought into scope only
-- after distribution, as the original body cannot reference it.
distributedLoopBody ::
  FlattenOps ->
  Segments ->
  SubExp ->
  Scope GPU ->
  DistInputs ->
  DistEnv ->
  [DistResult] ->
  Body SOACS ->
  FlattenM (Body GPU)
distributedLoopBody :: FlattenOps
-> [SubExp]
-> SubExp
-> Scope GPU
-> DistInputs
-> DistEnv
-> [DistResult]
-> Body SOACS
-> FlattenM (Body GPU)
distributedLoopBody FlattenOps
ops [SubExp]
segments SubExp
num_segments Scope GPU
loop_scope DistInputs
inputs DistEnv
env [DistResult]
res Body SOACS
body = do
  scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops
      (inputs', dstms) =
        distributeBodyWith ops scope segments inputs body
  buildBody_ $ localScope loop_scope $ do
    env' <- foldM (flattenDistStm ops segments) env dstms
    concat <$> zipWithM (liftLoopResult lvl segments num_segments inputs' env') res (bodyResult body)

-- | Make the original loop parameters available as distribution
-- inputs for the loop body, mapped to their lifted representations.
loopBodyInputs :: DistEnv -> DistInputs -> [FParam SOACS] -> [ResRep] -> FlattenM (DistInputs, DistEnv)
loopBodyInputs :: DistEnv
-> DistInputs
-> [FParam SOACS]
-> [ResRep]
-> FlattenM (DistInputs, DistEnv)
loopBodyInputs DistEnv
env DistInputs
inps [FParam SOACS]
old_loop_params [ResRep]
lifted_loop_reps = do
  (inps_local, env_local, next) <- DistEnv -> DistInputs -> FlattenM (DistInputs, DistEnv, Int)
localiseInputs DistEnv
env DistInputs
inps
  let loop_param_inputs =
        (Param DeclType -> Int -> (VName, DistInput))
-> [Param DeclType] -> [Int] -> DistInputs
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          (\Param DeclType
p Int
j -> (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p, ResTag -> Type -> DistInput
DistInput (Int -> ResTag
ResTag Int
j) (Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p)))
          [Param DeclType]
[FParam SOACS]
old_loop_params
          [Int
next ..]
      loop_param_reps =
        (Int -> ResRep -> (ResTag, ResRep))
-> [Int] -> [ResRep] -> [(ResTag, ResRep)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
j ResRep
rep -> (Int -> ResTag
ResTag Int
j, ResRep
rep)) [Int
next ..] [ResRep]
lifted_loop_reps
  (inps_local <> loop_param_inputs,) <$> insertRepsM loop_param_reps env_local

-- transform a for-loop with a variant iteration count into a while-loop
transformForToWhile ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  StmAux () ->
  [(FParam SOACS, SubExp)] ->
  VName ->
  IntType ->
  SubExp ->
  Body SOACS ->
  FlattenM DistEnv
transformForToWhile :: FlattenOps
-> [SubExp]
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> [(FParam SOACS, SubExp)]
-> VName
-> IntType
-> SubExp
-> Body SOACS
-> FlattenM DistEnv
transformForToWhile FlattenOps
ops [SubExp]
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux [(FParam SOACS, SubExp)]
merge VName
i IntType
it SubExp
n Body SOACS
body = do
  let old_loop_params :: [Param DeclType]
old_loop_params = ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
merge
  -- Fresh names used only in the synthetic rewritten body.
  cond_param_v <- Name -> FlattenM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"for_cond"
  cond0_v <- newVName "for_cond0"
  cond_next_v <- newVName "for_cond_next"
  i_next_v <- newVName "for_i_next"
  loop_old_out_vs <- replicateM (length merge) $ newVName "for_out"
  i_out_v <- newVName "for_i_out"
  cond_out_v <- newVName "for_cond_out"

  let zero = IntType -> Integer -> SubExp
intConst IntType
it Integer
0
      one = IntType -> Integer -> SubExp
intConst IntType
it Integer
1
      aux_no_certs = StmAux ()
aux {stmAuxCerts = mempty}

      cond0_stm =
        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 [VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
cond0_v (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool)])
          StmAux ()
StmAux (ExpDec SOACS)
aux_no_certs
          (BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSlt IntType
it) SubExp
zero SubExp
n)

      -- Extend the loop parameters with iteration variable and condition variable
      i_param = Attrs -> VName -> TypeBase shape o -> Param (TypeBase shape o)
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
i (PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim (IntType -> PrimType
IntType IntType
it))
      cond_param = Attrs -> VName -> TypeBase shape o -> Param (TypeBase shape o)
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
cond_param_v (PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool)

      Body loop_body_dec loop_body_stms loop_body_res = body

      i_next_stm =
        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 [VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
i_next_v (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim (IntType -> PrimType
IntType IntType
it))])
          StmAux ()
StmAux (ExpDec SOACS)
aux_no_certs
          -- OverflowWrap or OverflowUndef?
          (BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Overflow -> BinOp
Add IntType
it Overflow
OverflowUndef) (VName -> SubExp
Var VName
i) SubExp
one)

      cond_next_stm =
        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 [VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
cond_next_v (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool)])
          StmAux ()
StmAux (ExpDec SOACS)
aux_no_certs
          (BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSlt IntType
it) (VName -> SubExp
Var VName
i_next_v) SubExp
n)

      loop_new_body =
        BodyDec SOACS -> Stms SOACS -> [SubExpRes] -> Body SOACS
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body
          ()
BodyDec SOACS
loop_body_dec
          (Stms SOACS
loop_body_stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
i_next_stm Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
cond_next_stm)
          ( [ Certs -> SubExp -> SubExpRes
SubExpRes Certs
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
cond_next_v),
              Certs -> SubExp -> SubExpRes
SubExpRes Certs
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
i_next_v)
            ]
              [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [SubExpRes]
loop_body_res
          )

      merge' =
        [ (Param DeclType
forall {shape} {o}. Param (TypeBase shape o)
cond_param, VName -> SubExp
Var VName
cond0_v),
          (Param DeclType
forall {shape} {o}. Param (TypeBase shape o)
i_param, SubExp
zero)
        ]
          [(Param DeclType, SubExp)]
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. Semigroup a => a -> a -> a
<> [(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
merge

      loop_out_tys = [PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool, PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim (IntType -> PrimType
IntType IntType
it)] [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ (Param DeclType -> Type) -> [Param DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType [Param DeclType]
old_loop_params

      loop_pat =
        [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$
          (VName -> Type -> PatElem Type)
-> [VName] -> [Type] -> [PatElem Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
            VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem
            ([VName
cond_out_v, VName
i_out_v] [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
loop_old_out_vs)
            [Type]
loop_out_tys

      while_stm =
        Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let
          Pat Type
Pat (LetDec SOACS)
loop_pat
          StmAux ()
StmAux (ExpDec SOACS)
aux
          ([(FParam SOACS, SubExp)] -> LoopForm -> Body SOACS -> Exp SOACS
forall rep.
[(FParam rep, SubExp)] -> LoopForm -> Body rep -> Exp rep
Loop [(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
merge' (VName -> LoopForm
WhileLoop (Param (TypeBase (ZonkAny 0) (ZonkAny 1)) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase (ZonkAny 0) (ZonkAny 1))
forall {shape} {o}. Param (TypeBase shape o)
cond_param)) Body SOACS
loop_new_body)

      synthetic_body =
        BodyDec SOACS -> Stms SOACS -> [SubExpRes] -> Body SOACS
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body
          ()
          (Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
cond0_stm Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
while_stm)
          ((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]
loop_old_out_vs)

  reps <- distributeAndFlattenBody ops segments "for_variant_res" env inps res synthetic_body
  insertRepsM (zip (map distResTag res) reps) env

flattenLoop ::
  FlattenOps ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  (Pat Type, StmAux ()) ->
  ([(Param DeclType, SubExp)], LoopForm, Body SOACS) ->
  FlattenM DistEnv
flattenLoop :: FlattenOps
-> [SubExp]
-> DistEnv
-> DistInputs
-> [DistResult]
-> (Pat Type, StmAux ())
-> ([(Param DeclType, SubExp)], LoopForm, Body SOACS)
-> FlattenM DistEnv
flattenLoop FlattenOps
ops [SubExp]
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
_pat, StmAux ()
aux) ([(Param DeclType, SubExp)]
merge, ForLoop VName
i IntType
it SubExp
n, Body SOACS
body) = do
  if DistInputs -> SubExp -> Bool
isVariant DistInputs
inps SubExp
n
    then FlattenOps
-> [SubExp]
-> DistEnv
-> DistInputs
-> [DistResult]
-> StmAux ()
-> [(FParam SOACS, SubExp)]
-> VName
-> IntType
-> SubExp
-> Body SOACS
-> FlattenM DistEnv
transformForToWhile FlattenOps
ops [SubExp]
segments DistEnv
env DistInputs
inps [DistResult]
res StmAux ()
aux [(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
merge VName
i IntType
it SubExp
n Body SOACS
body
    else do
      let old_loop_params :: [Param DeclType]
old_loop_params = ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge
          loopParamNames :: Set VName
loopParamNames = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
paramName [Param DeclType]
old_loop_params

      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)
      (lifted_loop_params, lifted_loop_reps, lifted_init) <-
        unzip3 <$> mapM (liftLoopParam (flattenSegLevel ops) segments num_segments inps env loopParamNames) merge

      let lifted_loop_params' = [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
lifted_loop_params
          lifted_init' = [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[SubExp]]
lifted_init

      let i_param = Attrs -> VName -> TypeBase shape o -> Param (TypeBase shape o)
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
i (PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim (IntType -> PrimType
IntType IntType
it))
          build_scope = [FParam GPU] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param DeclType]
[FParam GPU]
lifted_loop_params' Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> [Param Type] -> Scope GPU
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param Type
forall {shape} {o}. Param (TypeBase shape o)
i_param]

      (loop_new_inputs, loop_env_local) <-
        localScope build_scope $
          loopBodyInputs env inps old_loop_params lifted_loop_reps

      -- When the loop parameters and all inputs used by the body are regular,
      -- the interchange of the map-nest and the loop corresponds to a perfectly
      -- ordinary Screma inside the loop. We then transform that Screma as if
      -- that was what the program looked like in the first place, which in
      -- particular means it is subject to multi-versioning. Otherwise we
      -- distribute the loop body statement by statement.
      let body_free = Body SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Body SOACS
body
          free_inps =
            [ (VName
v, DistInput
inp)
            | (VName
v, DistInput
inp) <- DistInputs
inps,
              VName
v VName -> Names -> Bool
`nameIn` Names
body_free,
              Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
loopParamNames
            ]
          regularInput (a
_, DistInput
inp) =
            Bool -> Bool
not ((SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (DistInputs -> Set VName -> SubExp -> Bool
variantDim DistInputs
inps Set VName
loopParamNames) (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (DistInput -> Type
distInputType DistInput
inp)))
              Bool -> Bool -> Bool
&& case DistInput
inp of
                DistInputFree {} -> Bool
True
                DistInput ResTag
rt Type
_ -> case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
                  Regular {} -> Bool
True
                  Irregular {} -> Bool
False
          regularRep Regular {} = Bool
True
          regularRep Irregular {} = Bool
False
          simpleParam t
p = case t -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf t
p of
            Prim {} -> Bool
True
            Array {} -> Bool
True
            DeclType
_ -> Bool
False
          -- The interchanged Screma is transformed as a top-level
          -- statement, so this is only possible when we are not
          -- generating in-block code.
          at_host_level = case FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops of
            SegThreadInBlock {} -> Bool
False
            SegLevel
_ -> Bool
True
          interchangeable =
            Bool
at_host_level
              -- Parameters with variant dimensions are lifted to an
              -- irregular representation, so this also rejects those.
              Bool -> Bool -> Bool
&& (ResRep -> Bool) -> [ResRep] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ResRep -> Bool
regularRep [ResRep]
lifted_loop_reps
              Bool -> Bool -> Bool
&& (DistResult -> Bool) -> [DistResult] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all DistResult -> Bool
isRegularDistResult [DistResult]
res
              Bool -> Bool -> Bool
&& (Param DeclType -> Bool) -> [Param DeclType] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Param DeclType -> Bool
forall {t}. DeclTyped t => t -> Bool
simpleParam [Param DeclType]
old_loop_params
              Bool -> Bool -> Bool
&& ((VName, DistInput) -> Bool) -> DistInputs -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName, DistInput) -> Bool
forall {a}. (a, DistInput) -> Bool
regularInput DistInputs
free_inps

      loop_body_gpu <-
        if interchangeable
          then
            localScope build_scope $
              interchangedLoopBody
                ops
                num_segments
                segments
                env
                (zip old_loop_params lifted_loop_params')
                free_inps
                aux
                body
          else
            distributedLoopBody
              ops
              segments
              num_segments
              build_scope
              loop_new_inputs
              loop_env_local
              res
              body

      let loop_exp_gpu =
            [(FParam GPU, SubExp)] -> LoopForm -> Body GPU -> Exp GPU
forall rep.
[(FParam rep, SubExp)] -> LoopForm -> Body rep -> Exp rep
Loop
              ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
lifted_loop_params' [SubExp]
lifted_init')
              (VName -> IntType -> SubExp -> LoopForm
ForLoop VName
i IntType
it SubExp
n)
              Body GPU
loop_body_gpu

      -- We must copy the result because otherwise we increase the degree of
      -- aliasing. In a loop, the result aliases the input, because it might run
      -- for zero iterations, but in the original program the result was
      -- produced by 'map', which has no aliases.
      loop_out_vs <-
        mapM (letExp "loop_res_out_copy" <=< eCopy . eVar)
          <=< certifying (distCerts inps aux env)
          $ letTupExp "loop_res_out" loop_exp_gpu

      let out_reps = [DistResult] -> [VName] -> [ResRep]
resultToResRepsByDistResult [DistResult]
res [VName]
loop_out_vs
      insertRepsM (zip (map distResTag res) out_reps) env
--
flattenLoop FlattenOps
ops [SubExp]
segments DistEnv
env DistInputs
inps [DistResult]
res (Pat Type
_pat, StmAux ()
aux) ([(Param DeclType, SubExp)]
merge, WhileLoop VName
cond, Body SOACS
body) = do
  -- TODO: Consider updating the active segment so we don't go over w every
  -- time.
  --
  -- inside the body we should compute the indices for which the condition is
  -- true and for which it is false, and then distribute the body based on that.
  -- We can then merge the results of the two branches by writing them back to a
  -- blank space like we do for the branches of a match.
  --
  -- This is probably not worth it: it is faster only for uniform loop
  -- parameters, but the common flattened-while case carries irregular state,
  -- which cannot be narrowed - an irregular result's offsets are global, so
  -- evicting a finished segment needs the final sizes of all segments, which
  -- are not known until the loop ends.

  let old_loop_params :: [Param DeclType]
old_loop_params = ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge
      loopParamNames :: Set VName
loopParamNames = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
paramName [Param DeclType]
old_loop_params
  w <- 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)
  (lifted_loop_params, lifted_loop_reps, lifted_init) <-
    unzip3 <$> mapM (liftLoopParam lvl segments w inps env loopParamNames) merge

  let lifted_loop_params' = [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
lifted_loop_params
      lifted_init' = [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[SubExp]]
lifted_init
      loop_params_scope = [FParam GPU] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param DeclType]
[FParam GPU]
lifted_loop_params'
  (loop_new_inputs, loop_env_local) <-
    localScope loop_params_scope $
      loopBodyInputs env inps old_loop_params lifted_loop_reps

  -- find cond_lifted_param in old_lifted_loop_params to get the lifted_loop_reps
  let maybe_cond = VName -> [(VName, (ResRep, [SubExp]))] -> Maybe (ResRep, [SubExp])
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
cond ([VName] -> [(ResRep, [SubExp])] -> [(VName, (ResRep, [SubExp]))]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
paramName [Param DeclType]
old_loop_params) ([ResRep] -> [[SubExp]] -> [(ResRep, [SubExp])]
forall a b. [a] -> [b] -> [(a, b)]
zip [ResRep]
lifted_loop_reps [[SubExp]]
lifted_init))
  scope <- askScope
  case maybe_cond of
    -- infinite loop
    Maybe (ResRep, [SubExp])
Nothing -> do
      loop_body_gpu <-
        FlattenOps
-> [SubExp]
-> SubExp
-> Scope GPU
-> DistInputs
-> DistEnv
-> [DistResult]
-> Body SOACS
-> FlattenM (Body GPU)
distributedLoopBody FlattenOps
ops [SubExp]
segments SubExp
w Scope GPU
loop_params_scope DistInputs
loop_new_inputs DistEnv
loop_env_local [DistResult]
res Body SOACS
body
      let loop_exp_gpu = [(FParam GPU, SubExp)] -> LoopForm -> Body GPU -> Exp GPU
forall rep.
[(FParam rep, SubExp)] -> LoopForm -> Body rep -> Exp rep
Loop ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
lifted_loop_params' [SubExp]
lifted_init') (VName -> LoopForm
WhileLoop VName
cond) Body GPU
loop_body_gpu
      loop_out_vs <- certifying (distCerts inps aux env) $ letTupExp "loop_res_out" loop_exp_gpu
      let out_reps = [DistResult] -> [VName] -> [ResRep]
resultToResRepsByDistResult [DistResult]
res [VName]
loop_out_vs
      insertRepsM (zip (map distResTag res) out_reps) env
    Just (ResRep
cond_lifted_rep, [SubExp]
cond_init) -> do
      let [SubExp
cond_init_se] = [SubExp]
cond_init

      -- Compute initial any_active
      cond_init_arr_v <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"cond_init_arr" (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 -> BasicOp
SubExp SubExp
cond_init_se
      let cond_lifted_param = case ResRep
cond_lifted_rep of
            Regular VName
v -> VName
v
            Irregular {} -> String -> VName
forall a. HasCallStack => String -> a
error String
"WhileLoop condition cannot be irregular"

      cond_init_arr_t <- lookupType cond_init_arr_v
      cond_init_flat <-
        letExp "cond_init_flat" . BasicOp $
          Reshape cond_init_arr_v $
            reshapeAll (arrayShape cond_init_arr_t) (Shape [w])

      or_lam <- binOpLambda LogOr Bool
      map_lam <- mkIdentityLambda [Prim Bool]
      ~[any_active_init_v] <-
        genNonSegRed lvl "any_active_init" [w] (Reduce Commutative or_lam [constant False]) mempty map_lam [cond_init_flat]
      let any_active_init = VName -> SubExp
Var VName
any_active_init_v

      any_active_param <- newParam "any_active" (Prim Bool)
      let build_scope = Scope GPU
loop_params_scope Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> [Param DeclType] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param DeclType
any_active_param]
      -- ‌build body
      loop_body_gpu <-
        buildBody_ . localScope build_scope $ do
          -- (num_data, active_inds) <- genFilter cond_lifted_param
          equiv_classes <- letExp "equiv_classes" <=< segMap lvl (MkSolo w) $ \(MkSolo SubExp
i) -> do
            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 ([SubExp] -> [TPrimExp Int64 VName]
segmentDims [SubExp]
segments) (SubExp -> TPrimExp Int64 VName
pe64 SubExp
i)
            c <- Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"c" (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
cond_lifted_param ((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)
            cls <-
              letSubExp "cls"
                =<< eIf
                  (eSubExp c)
                  (eBody [toExp $ intConst Int64 1])
                  (eBody [toExp $ intConst Int64 0])
            pure [subExpRes cls]
          n_cases <- letExp "n_cases" <=< toExp $ intConst Int64 2
          (partition_sizes, partition_offs, partition_inds) <- doPartition lvl n_cases equiv_classes
          inds_t <- lookupType partition_inds

          let getInds Name
nm Integer
k = do
                sz <-
                  Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp (Name
nm Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_sz")
                    (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 VName
partition_sizes [SubExp -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp (SubExp -> m (Exp (Rep m))) -> SubExp -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
k]
                off <-
                  letSubExp (nm <> "_off")
                    =<< eIndex partition_offs [toExp $ intConst Int64 k]
                inds <-
                  letExp (nm <> "_inds") $
                    BasicOp $
                      Index partition_inds $
                        fullSlice inds_t [DimSlice off sz (intConst Int64 1)]
                pure (sz, inds)

          (_, inactive_inds) <- getInds "inactive" 0
          (active_size, active_inds) <- getInds "active" 1

          inactive_reps <- forM old_loop_params $ \Param DeclType
p -> do
            (_, _, rep) <- SegLevel
-> [SubExp]
-> DistEnv
-> DistInputs
-> VName
-> Map VName ResRep
-> VName
-> FlattenM (Type, VName, ResRep)
splitInput SegLevel
lvl [SubExp]
segments DistEnv
loop_env_local DistInputs
loop_new_inputs VName
inactive_inds Map VName ResRep
forall a. Monoid a => a
mempty (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)
            pure rep

          let free_in_body =
                (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter
                  (DistInputs -> SubExp -> Bool
isVariant DistInputs
loop_new_inputs (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)
          free_sizes <-
            foldMap freeIn <$> mapM (lookupInputType loop_new_inputs) free_in_body
          let free_variant_sizes = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (DistInputs -> SubExp -> Bool
isVariant DistInputs
loop_new_inputs (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 loop_env_local loop_new_inputs active_inds mempty) free_size_vars
          let subset_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)
              env_subset = 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 subset_segments = [SubExp
active_size]
          let (subset_inputs', subset_dstms) =
                distributeBodyWith ops scope subset_segments subset_inputs body
          env_subset' <- foldM (flattenDistStm ops subset_segments) env_subset subset_dstms
          active_reps <-
            zipWithM
              (liftDistResultRep lvl subset_segments subset_inputs' env_subset')
              res
              (bodyResult body)

          let mergeOneLifted TypeBase Shape o
t ResRep
rep0 ResRep
rep1
                | TypeBase Shape o -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc TypeBase Shape o
t = do
                    let (Regular VName
acc_res) = ResRep
rep1
                    [SubExpRes] -> FlattenM [SubExpRes]
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Certs -> SubExp -> SubExpRes
SubExpRes Certs
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
acc_res)]
                | Bool
otherwise =
                    case (ResRep
rep0, ResRep
rep1) of
                      (Regular VName
x0, Regular VName
x1) -> do
                        let initial_shape :: Shape
initial_shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> TypeBase Shape o -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase Shape o
t
                        let final_shape :: Shape
final_shape = [SubExp] -> Shape
segmentsShape [SubExp]
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> TypeBase Shape o -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase Shape o
t
                        let pt :: PrimType
pt = TypeBase Shape o -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType TypeBase Shape o
t
                        space <- Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"blank" (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 (PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt Shape
initial_shape NoMode
NoMode)

                        out <-
                          foldM
                            (scatterRegular lvl)
                            space
                            [(inactive_inds, x0), (active_inds, x1)]

                        out_type <- arrayShape <$> lookupType out
                        out_reshaped <-
                          letExp "out_reshaped" . BasicOp $
                            Reshape out $
                              reshapeAll out_type final_shape

                        pure [SubExpRes mempty (Var out_reshaped)]
                      (Irregular IrregularRep
ir0, Irregular IrregularRep
ir1) -> do
                        segsSpace <-
                          Name -> Exp (Rep FlattenM) -> FlattenM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"blank_segs"
                            (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 (PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
int64 ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) NoMode
NoMode)

                        segs <-
                          foldM
                            (scatterRegular lvl)
                            segsSpace
                            [(inactive_inds, irregularS ir0), (active_inds, irregularS ir1)]

                        (_, offsets, num_data) <- exScanAndSum lvl segs

                        let pt = TypeBase Shape o -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType TypeBase Shape o
t
                        elemsSpace <-
                          letExp "blank_elems"
                            =<< eBlank (Array pt (Shape [num_data]) NoMode)

                        elems <-
                          foldM
                            (scatterIrregular lvl offsets)
                            elemsSpace
                            [(inactive_inds, ir0), (active_inds, ir1)]

                        flags <- genFlags lvl num_data offsets

                        pure
                          [ SubExpRes mempty num_data,
                            SubExpRes mempty (Var segs),
                            SubExpRes mempty (Var flags),
                            SubExpRes mempty (Var offsets),
                            SubExpRes mempty (Var elems)
                          ]
                      (ResRep, ResRep)
_ -> String -> FlattenM [SubExpRes]
forall a. HasCallStack => String -> a
error String
"mergeOneLifted: mismatched reps"

          merged_results <-
            concat
              <$> zipWithM
                (\Param DeclType
p (ResRep
r0, ResRep
r1) -> DeclType -> ResRep -> ResRep -> FlattenM [SubExpRes]
forall {o}.
TypeBase Shape o -> ResRep -> ResRep -> FlattenM [SubExpRes]
mergeOneLifted (Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf Param DeclType
p) ResRep
r0 ResRep
r1)
                old_loop_params
                (zip inactive_reps active_reps)

          -- we have one extra iteration but it is better than extra reduction in the loop body,
          any_active <-
            letSubExp "any_active"
              =<< eIf
                (toExp $ pe64 active_size .==. 0)
                (eBody [eSubExp $ constant False])
                (eBody [eSubExp $ constant True])

          pure $ merged_results ++ [SubExpRes mempty any_active]

      let merge' =
            [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip
              ([Param DeclType]
lifted_loop_params' [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType
any_active_param])
              ([SubExp]
lifted_init' [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp
any_active_init])
      loop_out_vs <-
        certifying (distCerts inps aux env) $
          letTupExp "loop_res_out" $
            Loop
              merge'
              (WhileLoop (paramName any_active_param))
              loop_body_gpu
      let loop_out_vs' = [VName] -> [VName]
forall a. HasCallStack => [a] -> [a]
L.init [VName]
loop_out_vs
      let out_reps = [DistResult] -> [VName] -> [ResRep]
resultToResRepsByDistResult [DistResult]
res [VName]
loop_out_vs'
      insertRepsM (zip (map distResTag res) out_reps) env
  where
    lvl :: SegLevel
lvl = FlattenOps -> SegLevel
flattenSegLevel FlattenOps
ops