{-# LANGUAGE TypeFamilies #-}

-- | General definitions for the flattening transformation.
--
-- Defines not just the core monads that are involved, but also the various
-- representations, except perhaps the ones that are completely local to another
-- module.
module Futhark.Pass.Flatten.Monad
  ( IrregularKind (..),
    IrregularRep (..),
    ResRep (..),
    DistEnv (..),
    FlattenOps (..),

    -- * Flattening monad
    FlattenM,
    FlattenState (..),
    runFlattenM,

    -- * Demands
    BuiltinFn (..),
    LiftMode (..),
    DemandFn (..),
    demandLifted,
    demandBuiltin,

    -- * Insertions
    insertRepM,
    insertRepsM,
    insertIrregularM,
    insertRegulars,

    -- * Various
    resVar,
    inputReps,
    segsAndElems,
  )
where

import Control.Monad
import Control.Monad.State
import Data.Bifunctor (bimap, second)
import Data.Foldable
import Data.Map qualified as M
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
import Futhark.IR.GPU
import Futhark.IR.SOACS (SOACS)
import Futhark.Pass.Flatten.Distribute
import Futhark.Tools
import Prelude hiding (div, rem)

-- | If true, 'sanityCheck' blocks are evaluated.
doSanityCheck :: Bool
doSanityCheck :: Bool
doSanityCheck = Bool
True

-- | Run a sanity-check that may verify invariants. The idea is that these can
-- be disabled without affecting the correctness of the pass, although there is
-- no constructive guarantee that no important effects take place in here.
sanityCheck :: (Monad m) => m () -> m ()
sanityCheck :: forall (m :: * -> *). Monad m => m () -> m ()
sanityCheck = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
doSanityCheck

-- Note [Representation of Flat Arrays]
--
-- This flattening implementation uses largely the nomenclature and
-- structure described by Cosmin Oancea. In particular, consider an
-- irregular array 'A' where
--
--   - A has 'n' segments (outermost dimension).
--
--   - A has element type 't'.
--
--   - A has a total of 'm' elements (where 'm' is divisible by 'n',
--     and may indeed be 'm').
--
-- Then A is represented by the following arrays:
--
--   - A_D : [m]t; the "data array".
--
--   - A_S : [n]i64; the "shape array" giving the number of scalar elements of each segment.
--
--   - A_F : [m]bool; the "flag array", indicating when an element begins a
--     new segment.
--
--   - A_O : [n]i64; the offset array, indicating for each segment
--     where it starts in the data (and flag) array.
--
--   - A_II1 : [m]i64; the "segment indices"; a mapping from element
--     index to index of the segment it belongs to.
--
--   - A_II2 : [m]i64; the "inner indices"; a mapping from element index
--     to index within its corresponding segment.
--
-- The arrays that are not the data array are collectively called the
-- "structure arrays". All of the structure arrays can be computed
-- from each other, but conceptually they all coexist.
--
-- Note that we only consider the *outer* dimension to be the
-- "segments". Also, 't' may actually be an array itself (although in
-- this case, the shape of 't' must be invariant to all parallel
-- dimensions). The inner structure is preserved through code, not
-- data. (Or in practice, ad-hoc auxiliary arrays produced by code.)
-- In Cosmin's notation, we maintain only the information for the
-- outermost dimension.
--
-- As an example, consider an irregular array
--
--   A = [ [], [ [1,2,3], [4], [], [5,6] ], [ [7], [], [8,9,10] ] ]
--
-- then
--
--   n = 3
--
--   m = 10
--
--   A_D = [1,2,3,4,5,6,7,8,9,10]
--
--   A_S = [0, 6, 4]
--
--   A_F = [T,F,F,F,F,F,T,F,F,F]
--
--   A_O = [0, 0, 6]
--
--   A_II1 = [1,1,1,1,1,1,2,2,2,2]
--
--   A_II2 = [0,0,0,1,3,3,0,2,2,2]

data IrregularKind
  = Dense
  | Replicated
  deriving (Int -> IrregularKind -> ShowS
[IrregularKind] -> ShowS
IrregularKind -> [Char]
(Int -> IrregularKind -> ShowS)
-> (IrregularKind -> [Char])
-> ([IrregularKind] -> ShowS)
-> Show IrregularKind
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> IrregularKind -> ShowS
showsPrec :: Int -> IrregularKind -> ShowS
$cshow :: IrregularKind -> [Char]
show :: IrregularKind -> [Char]
$cshowList :: [IrregularKind] -> ShowS
showList :: [IrregularKind] -> ShowS
Show, IrregularKind -> IrregularKind -> Bool
(IrregularKind -> IrregularKind -> Bool)
-> (IrregularKind -> IrregularKind -> Bool) -> Eq IrregularKind
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: IrregularKind -> IrregularKind -> Bool
== :: IrregularKind -> IrregularKind -> Bool
$c/= :: IrregularKind -> IrregularKind -> Bool
/= :: IrregularKind -> IrregularKind -> Bool
Eq)

data IrregularRep = IrregularRep
  { -- | Array of size of each segment, type @[]i64@.
    IrregularRep -> VName
irregularS :: VName,
    IrregularRep -> VName
irregularF :: VName,
    IrregularRep -> VName
irregularO :: VName,
    IrregularRep -> VName
irregularD :: VName,
    IrregularRep -> IrregularKind
irregularK :: IrregularKind
  }
  deriving (Int -> IrregularRep -> ShowS
[IrregularRep] -> ShowS
IrregularRep -> [Char]
(Int -> IrregularRep -> ShowS)
-> (IrregularRep -> [Char])
-> ([IrregularRep] -> ShowS)
-> Show IrregularRep
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> IrregularRep -> ShowS
showsPrec :: Int -> IrregularRep -> ShowS
$cshow :: IrregularRep -> [Char]
show :: IrregularRep -> [Char]
$cshowList :: [IrregularRep] -> ShowS
showList :: [IrregularRep] -> ShowS
Show)

data ResRep
  = -- | This variable is represented completely straightforwardly- if it is an
    -- array, it is a regular array.
    Regular VName
  | -- | The representation of an irregular array.
    Irregular IrregularRep
  deriving (Int -> ResRep -> ShowS
[ResRep] -> ShowS
ResRep -> [Char]
(Int -> ResRep -> ShowS)
-> (ResRep -> [Char]) -> ([ResRep] -> ShowS) -> Show ResRep
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ResRep -> ShowS
showsPrec :: Int -> ResRep -> ShowS
$cshow :: ResRep -> [Char]
show :: ResRep -> [Char]
$cshowList :: [ResRep] -> ShowS
showList :: [ResRep] -> ShowS
Show)

newtype DistEnv = DistEnv {DistEnv -> Map ResTag ResRep
distResMap :: M.Map ResTag ResRep}

insertRep :: ResTag -> ResRep -> DistEnv -> DistEnv
insertRep :: ResTag -> ResRep -> DistEnv -> DistEnv
insertRep ResTag
rt ResRep
rep DistEnv
env = DistEnv
env {distResMap = M.insert rt rep $ distResMap env}

insertRepM :: ResTag -> ResRep -> DistEnv -> FlattenM DistEnv
insertRepM :: ResTag -> ResRep -> DistEnv -> FlattenM DistEnv
insertRepM ResTag
rt ResRep
rep DistEnv
env = do
  FlattenM () -> FlattenM ()
forall (m :: * -> *). Monad m => m () -> m ()
sanityCheck (FlattenM () -> FlattenM ()) -> FlattenM () -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ do
    case ResRep
rep of
      Regular VName
_ -> () -> FlattenM ()
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Irregular (IrregularRep VName
shape VName
flags VName
offsets VName
data_ IrregularKind
_kind) -> do
        shape_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
shape
        flags_t <- lookupType flags
        data_t <- lookupType data_
        offsets_t <- lookupType offsets

        unless (arrayRank flags_t == 1 && elemType flags_t == Bool) $
          error $
            "Invalid flag array type: " <> prettyString flags_t
        unless (arrayRank offsets_t == 1 && elemType offsets_t == int64) $
          error $
            "Invalid offsets array type: " <> prettyString offsets_t
        unless (arrayRank shape_t == 1 && elemType shape_t == int64) $
          error $
            "Invalid shape array type: " <> prettyString shape_t
        when (arrayRank data_t /= 1) $
          error $
            "Invalid data array array: " <> prettyString data_t
  DistEnv -> FlattenM DistEnv
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DistEnv -> FlattenM DistEnv) -> DistEnv -> FlattenM DistEnv
forall a b. (a -> b) -> a -> b
$ ResTag -> ResRep -> DistEnv -> DistEnv
insertRep ResTag
rt ResRep
rep DistEnv
env

insertRepsM :: [(ResTag, ResRep)] -> DistEnv -> FlattenM DistEnv
insertRepsM :: [(ResTag, ResRep)] -> DistEnv -> FlattenM DistEnv
insertRepsM =
  (DistEnv -> [(ResTag, ResRep)] -> FlattenM DistEnv)
-> [(ResTag, ResRep)] -> DistEnv -> FlattenM DistEnv
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((DistEnv -> [(ResTag, ResRep)] -> FlattenM DistEnv)
 -> [(ResTag, ResRep)] -> DistEnv -> FlattenM DistEnv)
-> (DistEnv -> [(ResTag, ResRep)] -> FlattenM DistEnv)
-> [(ResTag, ResRep)]
-> DistEnv
-> FlattenM DistEnv
forall a b. (a -> b) -> a -> b
$ (DistEnv -> (ResTag, ResRep) -> FlattenM DistEnv)
-> DistEnv -> [(ResTag, ResRep)] -> FlattenM DistEnv
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (((ResTag, ResRep) -> DistEnv -> FlattenM DistEnv)
-> DistEnv -> (ResTag, ResRep) -> FlattenM DistEnv
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((ResTag, ResRep) -> DistEnv -> FlattenM DistEnv)
 -> DistEnv -> (ResTag, ResRep) -> FlattenM DistEnv)
-> ((ResTag, ResRep) -> DistEnv -> FlattenM DistEnv)
-> DistEnv
-> (ResTag, ResRep)
-> FlattenM DistEnv
forall a b. (a -> b) -> a -> b
$ (ResTag -> ResRep -> DistEnv -> FlattenM DistEnv)
-> (ResTag, ResRep) -> DistEnv -> FlattenM DistEnv
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ResTag -> ResRep -> DistEnv -> FlattenM DistEnv
insertRepM)

insertReps :: [(ResTag, ResRep)] -> DistEnv -> DistEnv
insertReps :: [(ResTag, ResRep)] -> DistEnv -> DistEnv
insertReps = (DistEnv -> [(ResTag, ResRep)] -> DistEnv)
-> [(ResTag, ResRep)] -> DistEnv -> DistEnv
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((DistEnv -> [(ResTag, ResRep)] -> DistEnv)
 -> [(ResTag, ResRep)] -> DistEnv -> DistEnv)
-> (DistEnv -> [(ResTag, ResRep)] -> DistEnv)
-> [(ResTag, ResRep)]
-> DistEnv
-> DistEnv
forall a b. (a -> b) -> a -> b
$ (DistEnv -> (ResTag, ResRep) -> DistEnv)
-> DistEnv -> [(ResTag, ResRep)] -> DistEnv
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (((ResTag, ResRep) -> DistEnv -> DistEnv)
-> DistEnv -> (ResTag, ResRep) -> DistEnv
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((ResTag, ResRep) -> DistEnv -> DistEnv)
 -> DistEnv -> (ResTag, ResRep) -> DistEnv)
-> ((ResTag, ResRep) -> DistEnv -> DistEnv)
-> DistEnv
-> (ResTag, ResRep)
-> DistEnv
forall a b. (a -> b) -> a -> b
$ (ResTag -> ResRep -> DistEnv -> DistEnv)
-> (ResTag, ResRep) -> DistEnv -> DistEnv
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ResTag -> ResRep -> DistEnv -> DistEnv
insertRep)

insertIrregularM :: VName -> VName -> VName -> ResTag -> VName -> IrregularKind -> DistEnv -> FlattenM DistEnv
insertIrregularM :: VName
-> VName
-> VName
-> ResTag
-> VName
-> IrregularKind
-> DistEnv
-> FlattenM DistEnv
insertIrregularM VName
shape VName
flags VName
offsets ResTag
rt VName
data_ IrregularKind
kind DistEnv
env = do
  let rep :: ResRep
rep = IrregularRep -> ResRep
Irregular (IrregularRep -> ResRep) -> IrregularRep -> ResRep
forall a b. (a -> b) -> a -> b
$ VName -> VName -> VName -> VName -> IrregularKind -> IrregularRep
IrregularRep VName
shape VName
flags VName
offsets VName
data_ IrregularKind
kind
  ResTag -> ResRep -> DistEnv -> FlattenM DistEnv
insertRepM ResTag
rt ResRep
rep DistEnv
env

insertRegulars :: [ResTag] -> [VName] -> DistEnv -> DistEnv
insertRegulars :: [ResTag] -> [VName] -> DistEnv -> DistEnv
insertRegulars [ResTag]
rts [VName]
xs =
  [(ResTag, ResRep)] -> DistEnv -> DistEnv
insertReps ([ResTag] -> [ResRep] -> [(ResTag, ResRep)]
forall a b. [a] -> [b] -> [(a, b)]
zip [ResTag]
rts ([ResRep] -> [(ResTag, ResRep)]) -> [ResRep] -> [(ResTag, ResRep)]
forall a b. (a -> b) -> a -> b
$ (VName -> ResRep) -> [VName] -> [ResRep]
forall a b. (a -> b) -> [a] -> [b]
map VName -> ResRep
Regular [VName]
xs)

instance Monoid DistEnv where
  mempty :: DistEnv
mempty = Map ResTag ResRep -> DistEnv
DistEnv Map ResTag ResRep
forall a. Monoid a => a
mempty

instance Semigroup DistEnv where
  DistEnv Map ResTag ResRep
x <> :: DistEnv -> DistEnv -> DistEnv
<> DistEnv Map ResTag ResRep
y = Map ResTag ResRep -> DistEnv
DistEnv (Map ResTag ResRep
x Map ResTag ResRep -> Map ResTag ResRep -> Map ResTag ResRep
forall a. Semigroup a => a -> a -> a
<> Map ResTag ResRep
y)

resVar :: ResTag -> DistEnv -> ResRep
resVar :: ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env = ResRep -> Maybe ResRep -> ResRep
forall a. a -> Maybe a -> a
fromMaybe ResRep
bad (Maybe ResRep -> ResRep) -> Maybe ResRep -> ResRep
forall a b. (a -> b) -> a -> b
$ ResTag -> Map ResTag ResRep -> Maybe ResRep
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ResTag
rt (Map ResTag ResRep -> Maybe ResRep)
-> Map ResTag ResRep -> Maybe ResRep
forall a b. (a -> b) -> a -> b
$ DistEnv -> Map ResTag ResRep
distResMap DistEnv
env
  where
    bad :: ResRep
bad = [Char] -> ResRep
forall a. HasCallStack => [Char] -> a
error ([Char] -> ResRep) -> [Char] -> ResRep
forall a b. (a -> b) -> a -> b
$ [Char]
"resVar: unknown tag: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ ResTag -> [Char]
forall a. Show a => a -> [Char]
show ResTag
rt

segsAndElems :: DistEnv -> [DistInput] -> (Maybe (VName, VName, VName), [VName])
segsAndElems :: DistEnv -> [DistInput] -> (Maybe (VName, VName, VName), [VName])
segsAndElems DistEnv
_ [] = (Maybe (VName, VName, VName)
forall a. Maybe a
Nothing, [])
segsAndElems DistEnv
env (DistInputFree VName
v Type
_ : [DistInput]
vs) =
  ([VName] -> [VName])
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:) ((Maybe (VName, VName, VName), [VName])
 -> (Maybe (VName, VName, VName), [VName]))
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall a b. (a -> b) -> a -> b
$ DistEnv -> [DistInput] -> (Maybe (VName, VName, VName), [VName])
segsAndElems DistEnv
env [DistInput]
vs
segsAndElems DistEnv
env (DistInput ResTag
rt Type
_ : [DistInput]
vs) =
  case ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env of
    Regular VName
v' ->
      ([VName] -> [VName])
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (VName
v' VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:) ((Maybe (VName, VName, VName), [VName])
 -> (Maybe (VName, VName, VName), [VName]))
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall a b. (a -> b) -> a -> b
$ DistEnv -> [DistInput] -> (Maybe (VName, VName, VName), [VName])
segsAndElems DistEnv
env [DistInput]
vs
    Irregular (IrregularRep VName
segments VName
flags VName
offsets VName
elems IrregularKind
k) -> do
      case IrregularKind
k of
        IrregularKind
Dense -> do
          (Maybe (VName, VName, VName) -> Maybe (VName, VName, VName))
-> ([VName] -> [VName])
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Maybe (VName, VName, VName)
-> Maybe (VName, VName, VName) -> Maybe (VName, VName, VName)
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus (Maybe (VName, VName, VName)
 -> Maybe (VName, VName, VName) -> Maybe (VName, VName, VName))
-> Maybe (VName, VName, VName)
-> Maybe (VName, VName, VName)
-> Maybe (VName, VName, VName)
forall a b. (a -> b) -> a -> b
$ (VName, VName, VName) -> Maybe (VName, VName, VName)
forall a. a -> Maybe a
Just (VName
segments, VName
flags, VName
offsets)) (VName
elems VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:) ((Maybe (VName, VName, VName), [VName])
 -> (Maybe (VName, VName, VName), [VName]))
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall a b. (a -> b) -> a -> b
$ DistEnv -> [DistInput] -> (Maybe (VName, VName, VName), [VName])
segsAndElems DistEnv
env [DistInput]
vs
        IrregularKind
Replicated ->
          ([VName] -> [VName])
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (VName
flags VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:) ((Maybe (VName, VName, VName), [VName])
 -> (Maybe (VName, VName, VName), [VName]))
-> (Maybe (VName, VName, VName), [VName])
-> (Maybe (VName, VName, VName), [VName])
forall a b. (a -> b) -> a -> b
$ DistEnv -> [DistInput] -> (Maybe (VName, VName, VName), [VName])
segsAndElems DistEnv
env [DistInput]
vs

-- | Mapping from original variable names to their distributed resreps.
inputReps :: DistInputs -> DistEnv -> M.Map VName (Type, ResRep)
inputReps :: DistInputs -> DistEnv -> Map VName (Type, ResRep)
inputReps DistInputs
inputs DistEnv
env = [(VName, (Type, ResRep))] -> Map VName (Type, ResRep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, (Type, ResRep))] -> Map VName (Type, ResRep))
-> [(VName, (Type, ResRep))] -> Map VName (Type, ResRep)
forall a b. (a -> b) -> a -> b
$ ((VName, DistInput) -> (VName, (Type, ResRep)))
-> DistInputs -> [(VName, (Type, ResRep))]
forall a b. (a -> b) -> [a] -> [b]
map ((DistInput -> (Type, ResRep))
-> (VName, DistInput) -> (VName, (Type, ResRep))
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second DistInput -> (Type, ResRep)
getRep) DistInputs
inputs
  where
    getRep :: DistInput -> (Type, ResRep)
getRep DistInput
di = case DistInput
di of
      DistInput ResTag
rt Type
t -> (Type
t, ResTag -> DistEnv -> ResRep
resVar ResTag
rt DistEnv
env)
      DistInputFree VName
v' Type
t -> (Type
t, VName -> ResRep
Regular VName
v')

-- | A representation of the different kinds of builtin functions we can
-- generate. This is used to only generate the ones we actually need for a given
-- program.
data BuiltinFn
  = BuiltinSegIota
  | BuiltinRepIota
  | BuiltinPrefixSum
  | BuiltinPartition
  deriving (BuiltinFn -> BuiltinFn -> Bool
(BuiltinFn -> BuiltinFn -> Bool)
-> (BuiltinFn -> BuiltinFn -> Bool) -> Eq BuiltinFn
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BuiltinFn -> BuiltinFn -> Bool
== :: BuiltinFn -> BuiltinFn -> Bool
$c/= :: BuiltinFn -> BuiltinFn -> Bool
/= :: BuiltinFn -> BuiltinFn -> Bool
Eq, Eq BuiltinFn
Eq BuiltinFn =>
(BuiltinFn -> BuiltinFn -> Ordering)
-> (BuiltinFn -> BuiltinFn -> Bool)
-> (BuiltinFn -> BuiltinFn -> Bool)
-> (BuiltinFn -> BuiltinFn -> Bool)
-> (BuiltinFn -> BuiltinFn -> Bool)
-> (BuiltinFn -> BuiltinFn -> BuiltinFn)
-> (BuiltinFn -> BuiltinFn -> BuiltinFn)
-> Ord BuiltinFn
BuiltinFn -> BuiltinFn -> Bool
BuiltinFn -> BuiltinFn -> Ordering
BuiltinFn -> BuiltinFn -> BuiltinFn
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: BuiltinFn -> BuiltinFn -> Ordering
compare :: BuiltinFn -> BuiltinFn -> Ordering
$c< :: BuiltinFn -> BuiltinFn -> Bool
< :: BuiltinFn -> BuiltinFn -> Bool
$c<= :: BuiltinFn -> BuiltinFn -> Bool
<= :: BuiltinFn -> BuiltinFn -> Bool
$c> :: BuiltinFn -> BuiltinFn -> Bool
> :: BuiltinFn -> BuiltinFn -> Bool
$c>= :: BuiltinFn -> BuiltinFn -> Bool
>= :: BuiltinFn -> BuiltinFn -> Bool
$cmax :: BuiltinFn -> BuiltinFn -> BuiltinFn
max :: BuiltinFn -> BuiltinFn -> BuiltinFn
$cmin :: BuiltinFn -> BuiltinFn -> BuiltinFn
min :: BuiltinFn -> BuiltinFn -> BuiltinFn
Ord, Int -> BuiltinFn -> ShowS
[BuiltinFn] -> ShowS
BuiltinFn -> [Char]
(Int -> BuiltinFn -> ShowS)
-> (BuiltinFn -> [Char])
-> ([BuiltinFn] -> ShowS)
-> Show BuiltinFn
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BuiltinFn -> ShowS
showsPrec :: Int -> BuiltinFn -> ShowS
$cshow :: BuiltinFn -> [Char]
show :: BuiltinFn -> [Char]
$cshowList :: [BuiltinFn] -> ShowS
showList :: [BuiltinFn] -> ShowS
Show)

data LiftMode
  = UniformLift
  | NonUniformLift
  deriving (LiftMode -> LiftMode -> Bool
(LiftMode -> LiftMode -> Bool)
-> (LiftMode -> LiftMode -> Bool) -> Eq LiftMode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LiftMode -> LiftMode -> Bool
== :: LiftMode -> LiftMode -> Bool
$c/= :: LiftMode -> LiftMode -> Bool
/= :: LiftMode -> LiftMode -> Bool
Eq, Eq LiftMode
Eq LiftMode =>
(LiftMode -> LiftMode -> Ordering)
-> (LiftMode -> LiftMode -> Bool)
-> (LiftMode -> LiftMode -> Bool)
-> (LiftMode -> LiftMode -> Bool)
-> (LiftMode -> LiftMode -> Bool)
-> (LiftMode -> LiftMode -> LiftMode)
-> (LiftMode -> LiftMode -> LiftMode)
-> Ord LiftMode
LiftMode -> LiftMode -> Bool
LiftMode -> LiftMode -> Ordering
LiftMode -> LiftMode -> LiftMode
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: LiftMode -> LiftMode -> Ordering
compare :: LiftMode -> LiftMode -> Ordering
$c< :: LiftMode -> LiftMode -> Bool
< :: LiftMode -> LiftMode -> Bool
$c<= :: LiftMode -> LiftMode -> Bool
<= :: LiftMode -> LiftMode -> Bool
$c> :: LiftMode -> LiftMode -> Bool
> :: LiftMode -> LiftMode -> Bool
$c>= :: LiftMode -> LiftMode -> Bool
>= :: LiftMode -> LiftMode -> Bool
$cmax :: LiftMode -> LiftMode -> LiftMode
max :: LiftMode -> LiftMode -> LiftMode
$cmin :: LiftMode -> LiftMode -> LiftMode
min :: LiftMode -> LiftMode -> LiftMode
Ord, Int -> LiftMode -> ShowS
[LiftMode] -> ShowS
LiftMode -> [Char]
(Int -> LiftMode -> ShowS)
-> (LiftMode -> [Char]) -> ([LiftMode] -> ShowS) -> Show LiftMode
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LiftMode -> ShowS
showsPrec :: Int -> LiftMode -> ShowS
$cshow :: LiftMode -> [Char]
show :: LiftMode -> [Char]
$cshowList :: [LiftMode] -> ShowS
showList :: [LiftMode] -> ShowS
Show)

-- | Indicate the need for a function to be generated. Instead of immediately
-- generating them ourselves, we collect requirements from multiple flattening
-- operations and satisfy them in their entirety.
data DemandFn
  = -- | We need this function to be lifted.
    DemandLifted Name LiftMode
  | DemandBuiltin BuiltinFn
  deriving (DemandFn -> DemandFn -> Bool
(DemandFn -> DemandFn -> Bool)
-> (DemandFn -> DemandFn -> Bool) -> Eq DemandFn
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DemandFn -> DemandFn -> Bool
== :: DemandFn -> DemandFn -> Bool
$c/= :: DemandFn -> DemandFn -> Bool
/= :: DemandFn -> DemandFn -> Bool
Eq, Eq DemandFn
Eq DemandFn =>
(DemandFn -> DemandFn -> Ordering)
-> (DemandFn -> DemandFn -> Bool)
-> (DemandFn -> DemandFn -> Bool)
-> (DemandFn -> DemandFn -> Bool)
-> (DemandFn -> DemandFn -> Bool)
-> (DemandFn -> DemandFn -> DemandFn)
-> (DemandFn -> DemandFn -> DemandFn)
-> Ord DemandFn
DemandFn -> DemandFn -> Bool
DemandFn -> DemandFn -> Ordering
DemandFn -> DemandFn -> DemandFn
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DemandFn -> DemandFn -> Ordering
compare :: DemandFn -> DemandFn -> Ordering
$c< :: DemandFn -> DemandFn -> Bool
< :: DemandFn -> DemandFn -> Bool
$c<= :: DemandFn -> DemandFn -> Bool
<= :: DemandFn -> DemandFn -> Bool
$c> :: DemandFn -> DemandFn -> Bool
> :: DemandFn -> DemandFn -> Bool
$c>= :: DemandFn -> DemandFn -> Bool
>= :: DemandFn -> DemandFn -> Bool
$cmax :: DemandFn -> DemandFn -> DemandFn
max :: DemandFn -> DemandFn -> DemandFn
$cmin :: DemandFn -> DemandFn -> DemandFn
min :: DemandFn -> DemandFn -> DemandFn
Ord, Int -> DemandFn -> ShowS
[DemandFn] -> ShowS
DemandFn -> [Char]
(Int -> DemandFn -> ShowS)
-> (DemandFn -> [Char]) -> ([DemandFn] -> ShowS) -> Show DemandFn
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DemandFn -> ShowS
showsPrec :: Int -> DemandFn -> ShowS
$cshow :: DemandFn -> [Char]
show :: DemandFn -> [Char]
$cshowList :: [DemandFn] -> ShowS
showList :: [DemandFn] -> ShowS
Show)

data FlattenState = FlattenState
  { -- In order to generate more stable threshold names, we keep track of the
    -- numbers used for thresholds separately from the ordinary name source.
    FlattenState -> Int
stateThresholdCounter :: Int,
    FlattenState -> VNameSource
stateNameSource :: VNameSource,
    -- A set of those functions that we have emitted calls to, and which will
    -- need to be generated.
    FlattenState -> Set DemandFn
stateDemandFns :: S.Set DemandFn
  }

newtype FlattenM a = FlattenM (BuilderT GPU (State FlattenState) a)
  deriving
    ( (forall a b. (a -> b) -> FlattenM a -> FlattenM b)
-> (forall a b. a -> FlattenM b -> FlattenM a) -> Functor FlattenM
forall a b. a -> FlattenM b -> FlattenM a
forall a b. (a -> b) -> FlattenM a -> FlattenM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> FlattenM a -> FlattenM b
fmap :: forall a b. (a -> b) -> FlattenM a -> FlattenM b
$c<$ :: forall a b. a -> FlattenM b -> FlattenM a
<$ :: forall a b. a -> FlattenM b -> FlattenM a
Functor,
      Functor FlattenM
Functor FlattenM =>
(forall a. a -> FlattenM a)
-> (forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b)
-> (forall a b c.
    (a -> b -> c) -> FlattenM a -> FlattenM b -> FlattenM c)
-> (forall a b. FlattenM a -> FlattenM b -> FlattenM b)
-> (forall a b. FlattenM a -> FlattenM b -> FlattenM a)
-> Applicative FlattenM
forall a. a -> FlattenM a
forall a b. FlattenM a -> FlattenM b -> FlattenM a
forall a b. FlattenM a -> FlattenM b -> FlattenM b
forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b
forall a b c.
(a -> b -> c) -> FlattenM a -> FlattenM b -> FlattenM c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> FlattenM a
pure :: forall a. a -> FlattenM a
$c<*> :: forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b
<*> :: forall a b. FlattenM (a -> b) -> FlattenM a -> FlattenM b
$cliftA2 :: forall a b c.
(a -> b -> c) -> FlattenM a -> FlattenM b -> FlattenM c
liftA2 :: forall a b c.
(a -> b -> c) -> FlattenM a -> FlattenM b -> FlattenM c
$c*> :: forall a b. FlattenM a -> FlattenM b -> FlattenM b
*> :: forall a b. FlattenM a -> FlattenM b -> FlattenM b
$c<* :: forall a b. FlattenM a -> FlattenM b -> FlattenM a
<* :: forall a b. FlattenM a -> FlattenM b -> FlattenM a
Applicative,
      Applicative FlattenM
Applicative FlattenM =>
(forall a b. FlattenM a -> (a -> FlattenM b) -> FlattenM b)
-> (forall a b. FlattenM a -> FlattenM b -> FlattenM b)
-> (forall a. a -> FlattenM a)
-> Monad FlattenM
forall a. a -> FlattenM a
forall a b. FlattenM a -> FlattenM b -> FlattenM b
forall a b. FlattenM a -> (a -> FlattenM b) -> FlattenM b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. FlattenM a -> (a -> FlattenM b) -> FlattenM b
>>= :: forall a b. FlattenM a -> (a -> FlattenM b) -> FlattenM b
$c>> :: forall a b. FlattenM a -> FlattenM b -> FlattenM b
>> :: forall a b. FlattenM a -> FlattenM b -> FlattenM b
$creturn :: forall a. a -> FlattenM a
return :: forall a. a -> FlattenM a
Monad,
      LocalScope GPU,
      HasScope GPU,
      MonadState FlattenState,
      Monad FlattenM
FlattenM VNameSource
Monad FlattenM =>
FlattenM VNameSource
-> (VNameSource -> FlattenM ()) -> MonadFreshNames FlattenM
VNameSource -> FlattenM ()
forall (m :: * -> *).
Monad m =>
m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
$cgetNameSource :: FlattenM VNameSource
getNameSource :: FlattenM VNameSource
$cputNameSource :: VNameSource -> FlattenM ()
putNameSource :: VNameSource -> FlattenM ()
MonadFreshNames
    )

instance MonadBuilder FlattenM where
  type Rep FlattenM = GPU
  mkExpDecM :: Pat (LetDec (Rep FlattenM))
-> Exp (Rep FlattenM) -> FlattenM (ExpDec (Rep FlattenM))
mkExpDecM Pat (LetDec (Rep FlattenM))
pat Exp (Rep FlattenM)
e = BuilderT GPU (State FlattenState) (ExpDec (Rep FlattenM))
-> FlattenM (ExpDec (Rep FlattenM))
forall a. BuilderT GPU (State FlattenState) a -> FlattenM a
FlattenM (BuilderT GPU (State FlattenState) (ExpDec (Rep FlattenM))
 -> FlattenM (ExpDec (Rep FlattenM)))
-> BuilderT GPU (State FlattenState) (ExpDec (Rep FlattenM))
-> FlattenM (ExpDec (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (BuilderT GPU (State FlattenState))))
-> Exp (Rep (BuilderT GPU (State FlattenState)))
-> BuilderT
     GPU
     (State FlattenState)
     (ExpDec (Rep (BuilderT GPU (State FlattenState))))
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))
mkExpDecM Pat (LetDec (Rep (BuilderT GPU (State FlattenState))))
Pat (LetDec (Rep FlattenM))
pat Exp (Rep (BuilderT GPU (State FlattenState)))
Exp (Rep FlattenM)
e
  mkBodyM :: forall res.
IsResult res =>
Stms (Rep FlattenM) -> [res] -> FlattenM (GBody (Rep FlattenM) res)
mkBodyM Stms (Rep FlattenM)
stms [res]
res = BuilderT GPU (State FlattenState) (GBody (Rep FlattenM) res)
-> FlattenM (GBody (Rep FlattenM) res)
forall a. BuilderT GPU (State FlattenState) a -> FlattenM a
FlattenM (BuilderT GPU (State FlattenState) (GBody (Rep FlattenM) res)
 -> FlattenM (GBody (Rep FlattenM) res))
-> BuilderT GPU (State FlattenState) (GBody (Rep FlattenM) res)
-> FlattenM (GBody (Rep FlattenM) res)
forall a b. (a -> b) -> a -> b
$ Stms (Rep (BuilderT GPU (State FlattenState)))
-> [res]
-> BuilderT
     GPU
     (State FlattenState)
     (GBody (Rep (BuilderT GPU (State FlattenState))) res)
forall res.
IsResult res =>
Stms (Rep (BuilderT GPU (State FlattenState)))
-> [res]
-> BuilderT
     GPU
     (State FlattenState)
     (GBody (Rep (BuilderT GPU (State FlattenState))) res)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
Stms (Rep m) -> [res] -> m (GBody (Rep m) res)
mkBodyM Stms (Rep (BuilderT GPU (State FlattenState)))
Stms (Rep FlattenM)
stms [res]
res
  mkLetNamesM :: [VName] -> Exp (Rep FlattenM) -> FlattenM (Stm (Rep FlattenM))
mkLetNamesM [VName]
pat Exp (Rep FlattenM)
e = BuilderT GPU (State FlattenState) (Stm (Rep FlattenM))
-> FlattenM (Stm (Rep FlattenM))
forall a. BuilderT GPU (State FlattenState) a -> FlattenM a
FlattenM (BuilderT GPU (State FlattenState) (Stm (Rep FlattenM))
 -> FlattenM (Stm (Rep FlattenM)))
-> BuilderT GPU (State FlattenState) (Stm (Rep FlattenM))
-> FlattenM (Stm (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp (Rep (BuilderT GPU (State FlattenState)))
-> BuilderT
     GPU
     (State FlattenState)
     (Stm (Rep (BuilderT GPU (State FlattenState))))
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesM [VName]
pat Exp (Rep (BuilderT GPU (State FlattenState)))
Exp (Rep FlattenM)
e

  addStms :: Stms (Rep FlattenM) -> FlattenM ()
addStms = BuilderT GPU (State FlattenState) () -> FlattenM ()
forall a. BuilderT GPU (State FlattenState) a -> FlattenM a
FlattenM (BuilderT GPU (State FlattenState) () -> FlattenM ())
-> (Stms GPU -> BuilderT GPU (State FlattenState) ())
-> Stms GPU
-> FlattenM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms (Rep (BuilderT GPU (State FlattenState)))
-> BuilderT GPU (State FlattenState) ()
Stms GPU -> BuilderT GPU (State FlattenState) ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms
  collectStms :: forall a. FlattenM a -> FlattenM (a, Stms (Rep FlattenM))
collectStms (FlattenM BuilderT GPU (State FlattenState) a
m) = BuilderT GPU (State FlattenState) (a, Stms (Rep FlattenM))
-> FlattenM (a, Stms (Rep FlattenM))
forall a. BuilderT GPU (State FlattenState) a -> FlattenM a
FlattenM (BuilderT GPU (State FlattenState) (a, Stms (Rep FlattenM))
 -> FlattenM (a, Stms (Rep FlattenM)))
-> BuilderT GPU (State FlattenState) (a, Stms (Rep FlattenM))
-> FlattenM (a, Stms (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ BuilderT GPU (State FlattenState) a
-> BuilderT
     GPU
     (State FlattenState)
     (a, Stms (Rep (BuilderT GPU (State FlattenState))))
forall a.
BuilderT GPU (State FlattenState) a
-> BuilderT
     GPU
     (State FlattenState)
     (a, Stms (Rep (BuilderT GPU (State FlattenState))))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms BuilderT GPU (State FlattenState) a
m

instance MonadFreshNames (State FlattenState) where
  getNameSource :: State FlattenState VNameSource
getNameSource = (FlattenState -> VNameSource) -> State FlattenState VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets FlattenState -> VNameSource
stateNameSource
  putNameSource :: VNameSource -> State FlattenState ()
putNameSource VNameSource
src = (FlattenState -> FlattenState) -> State FlattenState ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FlattenState -> FlattenState) -> State FlattenState ())
-> (FlattenState -> FlattenState) -> State FlattenState ()
forall a b. (a -> b) -> a -> b
$ \FlattenState
s -> FlattenState
s {stateNameSource = src}

-- | Do not nest these - the counter for thresholds will be wrong.
runFlattenM :: (MonadFreshNames m) => Scope GPU -> FlattenM a -> m (a, S.Set DemandFn)
runFlattenM :: forall (m :: * -> *) a.
MonadFreshNames m =>
Scope GPU -> FlattenM a -> m (a, Set DemandFn)
runFlattenM Scope GPU
scope (FlattenM BuilderT GPU (State FlattenState) a
m) = (VNameSource -> ((a, Set DemandFn), VNameSource))
-> m (a, Set DemandFn)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ((a, Set DemandFn), VNameSource))
 -> m (a, Set DemandFn))
-> (VNameSource -> ((a, Set DemandFn), VNameSource))
-> m (a, Set DemandFn)
forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
  let initial_state :: FlattenState
initial_state =
        FlattenState
          { stateThresholdCounter :: Int
stateThresholdCounter = Int
0,
            stateNameSource :: VNameSource
stateNameSource = VNameSource
src,
            stateDemandFns :: Set DemandFn
stateDemandFns = Set DemandFn
forall a. Monoid a => a
mempty
          }
      (a
x, FlattenState
s) = State FlattenState a -> FlattenState -> (a, FlattenState)
forall s a. State s a -> s -> (a, s)
runState ((a, Stms GPU) -> a
forall a b. (a, b) -> a
fst ((a, Stms GPU) -> a)
-> StateT FlattenState Identity (a, Stms GPU)
-> State FlattenState a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BuilderT GPU (State FlattenState) a
-> Scope GPU -> StateT FlattenState Identity (a, Stms GPU)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT BuilderT GPU (State FlattenState) a
m Scope GPU
scope) FlattenState
initial_state
   in ((a
x, FlattenState -> Set DemandFn
stateDemandFns FlattenState
s), FlattenState -> VNameSource
stateNameSource FlattenState
s)

-- | Indicate that we rather need a lifted version of this function.
demandLifted :: Name -> LiftMode -> FlattenM ()
demandLifted :: Name -> LiftMode -> FlattenM ()
demandLifted Name
fname LiftMode
mode = (FlattenState -> FlattenState) -> FlattenM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FlattenState -> FlattenState) -> FlattenM ())
-> (FlattenState -> FlattenState) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ \FlattenState
s ->
  FlattenState
s {stateDemandFns = S.insert (DemandLifted fname mode) $ stateDemandFns s}

-- | Demand the presence of this builtin function.
demandBuiltin :: BuiltinFn -> FlattenM ()
demandBuiltin :: BuiltinFn -> FlattenM ()
demandBuiltin BuiltinFn
b = (FlattenState -> FlattenState) -> FlattenM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FlattenState -> FlattenState) -> FlattenM ())
-> (FlattenState -> FlattenState) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ \FlattenState
s ->
  FlattenState
s {stateDemandFns = S.insert (DemandBuiltin b) $ stateDemandFns s}

-- | Functions for tying together disparate modules - this is to avoid mutually
-- recursive modules.
data FlattenOps = FlattenOps
  { FlattenOps -> SegLevel
flattenSegLevel :: SegLevel,
    -- | How to treat irregularity when distributing the bodies we encounter.
    FlattenOps -> DistIrregularity
flattenIrregularity :: DistIrregularity,
    FlattenOps -> FunHasParallelism
flattenFunHasParallelism :: FunHasParallelism,
    -- | Flatten a 'DistStm' using the given ops, which need not be the ones
    -- this record belongs to - see 'atSegLevel'. Use 'flattenDistStm' to
    -- continue with the current ops.
    FlattenOps
-> FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv
flattenDistStmWith :: FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv,
    -- | Flatten a scalar statement at the given seg level, which need not be
    -- the one this record carries. Use 'flattenScalarStm' for the latter.
    FlattenOps
-> SegLevel
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Stm SOACS
-> FlattenM DistEnv
flattenScalarStmAt :: SegLevel -> Segments -> DistEnv -> DistInputs -> [DistResult] -> Stm SOACS -> FlattenM DistEnv,
    -- | Transform a statement as if it occurred at the top level, including
    -- multi-versioning of SOACs. Used when a transformation (e.g. loop
    -- interchange) produces a statement that should be treated as if the
    -- program had looked like that all along.
    FlattenOps -> Stm SOACS -> FlattenM ()
flattenTopLevelStm :: Stm SOACS -> FlattenM ()
  }