    {-# LANGUAGE LambdaCase #-}
    {-# LANGUAGE TypeFamilies #-}

    module Futhark.Pass.ExtractKernels.Distribution
      ( Target,
        Targets,
        ppTargets,
        singleTarget,
        outerTarget,
        innerTarget,
        pushInnerTarget,
        popInnerTarget,
        targetsScope,
        LoopNesting (..),
        ppLoopNesting,
        scopeOfLoopNesting,
        Nesting (..),
        Nestings,
        ppNestings,
        letBindInInnerNesting,
        singleNesting,
        pushInnerNesting,
        KernelNest,
        ppKernelNest,
        newKernel,
        innermostKernelNesting,
        pushKernelNesting,
        pushInnerKernelNesting,
        scopeOfKernelNest,
        kernelNestLoops,
        kernelNestWidths,
        boundInKernelNest,
        boundInKernelNests,
        flatKernel,
        constructKernel,
        tryDistribute,
        tryDistributeStm,
      )
    where

    import Control.Monad
    import Control.Monad.RWS.Strict
    import Control.Monad.Trans.Maybe
    import Data.Bifunctor (second)
    import Data.Foldable
    import Data.List (elemIndex, sortOn)
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Futhark.IR
    import Futhark.IR.SegOp
    import Futhark.MonadFreshNames
    import Futhark.Pass.ExtractKernels.BlockedKernel
      ( DistRep,
        KernelInput (..),
        MkSegLevel,
        mapKernel,
        readKernelInput,
      )
    import Futhark.Tools
    import Futhark.Transform.Rename
    import Futhark.Util
    import Futhark.Util.Log

    type Target = (Pat Type, Result)

    -- | First pair element is the very innermost ("current") target.  In
    -- the list, the outermost target comes first.  Invariant: Every
    -- element of a pattern must be present as the result of the
    -- immediately enclosing target.  This is ensured by 'pushInnerTarget'
    -- by removing unused pattern elements.
    data Targets = Targets
      { Targets -> Target
    _innerTarget :: Target,
        Targets -> [Target]
    _outerTargets :: [Target]
      }

    ppTargets :: Targets -> String
    ppTargets :: Targets -> [Char]
    ppTargets (Targets Target
    target [Target]
    targets) =
      [[Char]] -> [Char]
    unlines ([[Char]] -> [Char]) -> [[Char]] -> [Char]
    forall a b. (a -> b) -> a -> b
    $ (Target -> [Char]) -> [Target] -> [[Char]]
    forall a b. (a -> b) -> [a] -> [b]
    map Target -> [Char]
    forall {a} {a}. (Pretty a, Pretty a) => (a, a) -> [Char]
    ppTarget ([Target] -> [[Char]]) -> [Target] -> [[Char]]
    forall a b. (a -> b) -> a -> b
    $ [Target]
    targets [Target] -> [Target] -> [Target]
    forall a. [a] -> [a] -> [a]
    ++ [Target
    target]
      where
        ppTarget :: (a, a) -> [Char]
    ppTarget (a
    pat, a
    res) = a -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString a
    pat [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " <- " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ a -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString a
    res

    singleTarget :: Target -> Targets
    singleTarget :: Target -> Targets
    singleTarget = (Target -> [Target] -> Targets) -> [Target] -> Target -> Targets
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Target -> [Target] -> Targets
    Targets []

    outerTarget :: Targets -> Target
    outerTarget :: Targets -> Target
    outerTarget (Targets Target
    inner_target []) = Target
    inner_target
    outerTarget (Targets Target
    _ (Target
    outer_target : [Target]
    _)) = Target
    outer_target

    innerTarget :: Targets -> Target
    innerTarget :: Targets -> Target
    innerTarget (Targets Target
    inner_target [Target]
    _) = Target
    inner_target

    pushOuterTarget :: Target -> Targets -> Targets
    pushOuterTarget :: Target -> Targets -> Targets
    pushOuterTarget Target
    target (Targets Target
    inner_target [Target]
    targets) =
      Target -> [Target] -> Targets
    Targets Target
    inner_target (Target
    target Target -> [Target] -> [Target]
    forall a. a -> [a] -> [a]
    : [Target]
    targets)

    pushInnerTarget :: Target -> Targets -> Targets
    pushInnerTarget :: Target -> Targets -> Targets
    pushInnerTarget (Pat Type
    pat, Result
    res) (Targets Target
    inner_target [Target]
    targets) =
      Target -> [Target] -> Targets
    Targets (Pat Type
    pat', Result
    res') ([Target]
    targets [Target] -> [Target] -> [Target]
    forall a. [a] -> [a] -> [a]
    ++ [Target
    inner_target])
      where
        ([PatElem Type]
    pes', Result
    res') = [(PatElem Type, SubExpRes)] -> ([PatElem Type], Result)
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(PatElem Type, SubExpRes)] -> ([PatElem Type], Result))
    -> [(PatElem Type, SubExpRes)] -> ([PatElem Type], Result)
    forall a b. (a -> b) -> a -> b
    $ ((PatElem Type, SubExpRes) -> Bool)
    -> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (PatElem Type -> Bool
    used (PatElem Type -> Bool)
    -> ((PatElem Type, SubExpRes) -> PatElem Type)
    -> (PatElem Type, SubExpRes)
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (PatElem Type, SubExpRes) -> PatElem Type
    forall a b. (a, b) -> a
    fst) ([(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)])
    -> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
    forall a b. (a -> b) -> a -> b
    $ [PatElem Type] -> Result -> [(PatElem Type, SubExpRes)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    pat) Result
    res
        pat' :: Pat Type
    pat' = [PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [PatElem Type]
    pes'
        inner_used :: Names
    inner_used = Result -> Names
    forall a. FreeIn a => a -> Names
    freeIn (Result -> Names) -> Result -> Names
    forall a b. (a -> b) -> a -> b
    $ Target -> Result
    forall a b. (a, b) -> b
    snd Target
    inner_target
        used :: PatElem Type -> Bool
    used PatElem Type
    pe = PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem Type
    pe VName -> Names -> Bool
    `nameIn` Names
    inner_used

    popInnerTarget :: Targets -> Maybe (Target, Targets)
    popInnerTarget :: Targets -> Maybe (Target, Targets)
    popInnerTarget (Targets Target
    t [Target]
    ts) =
      case [Target] -> [Target]
    forall a. [a] -> [a]
    reverse [Target]
    ts of
        Target
    x : [Target]
    xs -> (Target, Targets) -> Maybe (Target, Targets)
    forall a. a -> Maybe a
    Just (Target
    t, Target -> [Target] -> Targets
    Targets Target
    x ([Target] -> Targets) -> [Target] -> Targets
    forall a b. (a -> b) -> a -> b
    $ [Target] -> [Target]
    forall a. [a] -> [a]
    reverse [Target]
    xs)
        [] -> Maybe (Target, Targets)
    forall a. Maybe a
    Nothing

    targetScope :: (DistRep rep) => Target -> Scope rep
    targetScope :: forall rep. DistRep rep => Target -> Scope rep
    targetScope = Pat Type -> Scope rep
    forall rep dec. (LetDec rep ~ dec) => Pat dec -> Scope rep
    scopeOfPat (Pat Type -> Scope rep)
    -> (Target -> Pat Type) -> Target -> Scope rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Target -> Pat Type
    forall a b. (a, b) -> a
    fst

    targetsScope :: (DistRep rep) => Targets -> Scope rep
    targetsScope :: forall rep. DistRep rep => Targets -> Scope rep
    targetsScope (Targets Target
    t [Target]
    ts) = [Scope rep] -> Scope rep
    forall a. Monoid a => [a] -> a
    mconcat ([Scope rep] -> Scope rep) -> [Scope rep] -> Scope rep
    forall a b. (a -> b) -> a -> b
    $ (Target -> Scope rep) -> [Target] -> [Scope rep]
    forall a b. (a -> b) -> [a] -> [b]
    map Target -> Scope rep
    forall rep. DistRep rep => Target -> Scope rep
    targetScope ([Target] -> [Scope rep]) -> [Target] -> [Scope rep]
    forall a b. (a -> b) -> a -> b
    $ Target
    t Target -> [Target] -> [Target]
    forall a. a -> [a] -> [a]
    : [Target]
    ts

    data LoopNesting = MapNesting
      { LoopNesting -> Pat Type
    loopNestingPat :: Pat Type,
        LoopNesting -> StmAux ()
    loopNestingAux :: StmAux (),
        LoopNesting -> SubExp
    loopNestingWidth :: SubExp,
        LoopNesting -> [(Param Type, VName)]
    loopNestingParamsAndArrs :: [(Param Type, VName)]
      }
      deriving (Int -> LoopNesting -> [Char] -> [Char]
    [LoopNesting] -> [Char] -> [Char]
    LoopNesting -> [Char]
    (Int -> LoopNesting -> [Char] -> [Char])
    -> (LoopNesting -> [Char])
    -> ([LoopNesting] -> [Char] -> [Char])
    -> Show LoopNesting
    forall a.
    (Int -> a -> [Char] -> [Char])
    -> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
    $cshowsPrec :: Int -> LoopNesting -> [Char] -> [Char]
    showsPrec :: Int -> LoopNesting -> [Char] -> [Char]
    $cshow :: LoopNesting -> [Char]
    show :: LoopNesting -> [Char]
    $cshowList :: [LoopNesting] -> [Char] -> [Char]
    showList :: [LoopNesting] -> [Char] -> [Char]
    Show)

    scopeOfLoopNesting :: (LParamInfo rep ~ Type) => LoopNesting -> Scope rep
    scopeOfLoopNesting :: forall rep. (LParamInfo rep ~ Type) => LoopNesting -> Scope rep
    scopeOfLoopNesting = [Param Type] -> Scope rep
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams ([Param Type] -> Scope rep)
    -> (LoopNesting -> [Param Type]) -> LoopNesting -> Scope rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((Param Type, VName) -> Param Type)
    -> [(Param Type, VName)] -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param Type, VName) -> Param Type
    forall a b. (a, b) -> a
    fst ([(Param Type, VName)] -> [Param Type])
    -> (LoopNesting -> [(Param Type, VName)])
    -> LoopNesting
    -> [Param Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . LoopNesting -> [(Param Type, VName)]
    loopNestingParamsAndArrs

    ppLoopNesting :: LoopNesting -> String
    ppLoopNesting :: LoopNesting -> [Char]
    ppLoopNesting (MapNesting Pat Type
    _ StmAux ()
    _ SubExp
    _ [(Param Type, VName)]
    params_and_arrs) =
      [Param Type] -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString (((Param Type, VName) -> Param Type)
    -> [(Param Type, VName)] -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param Type, VName) -> Param Type
    forall a b. (a, b) -> a
    fst [(Param Type, VName)]
    params_and_arrs)
        [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " <- "
        [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [VName] -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString (((Param Type, VName) -> VName) -> [(Param Type, VName)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param Type, VName) -> VName
    forall a b. (a, b) -> b
    snd [(Param Type, VName)]
    params_and_arrs)

    loopNestingParams :: LoopNesting -> [Param Type]
    loopNestingParams :: LoopNesting -> [Param Type]
    loopNestingParams = ((Param Type, VName) -> Param Type)
    -> [(Param Type, VName)] -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param Type, VName) -> Param Type
    forall a b. (a, b) -> a
    fst ([(Param Type, VName)] -> [Param Type])
    -> (LoopNesting -> [(Param Type, VName)])
    -> LoopNesting
    -> [Param Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . LoopNesting -> [(Param Type, VName)]
    loopNestingParamsAndArrs

    instance FreeIn LoopNesting where
      freeIn' :: LoopNesting -> FV
    freeIn' (MapNesting Pat Type
    pat StmAux ()
    aux SubExp
    w [(Param Type, VName)]
    params_and_arrs) =
        Pat Type -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Pat Type
    pat FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> StmAux () -> FV
    forall a. FreeIn a => a -> FV
    freeIn' StmAux ()
    aux FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SubExp
    w FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> [(Param Type, VName)] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [(Param Type, VName)]
    params_and_arrs

    data Nesting = Nesting
      { Nesting -> Names
    nestingLetBound :: Names,
        Nesting -> LoopNesting
    nestingLoop :: LoopNesting
      }
      deriving (Int -> Nesting -> [Char] -> [Char]
    [Nesting] -> [Char] -> [Char]
    Nesting -> [Char]
    (Int -> Nesting -> [Char] -> [Char])
    -> (Nesting -> [Char])
    -> ([Nesting] -> [Char] -> [Char])
    -> Show Nesting
    forall a.
    (Int -> a -> [Char] -> [Char])
    -> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
    $cshowsPrec :: Int -> Nesting -> [Char] -> [Char]
    showsPrec :: Int -> Nesting -> [Char] -> [Char]
    $cshow :: Nesting -> [Char]
    show :: Nesting -> [Char]
    $cshowList :: [Nesting] -> [Char] -> [Char]
    showList :: [Nesting] -> [Char] -> [Char]
    Show)

    letBindInNesting :: Names -> Nesting -> Nesting
    letBindInNesting :: Names -> Nesting -> Nesting
    letBindInNesting Names
    newnames (Nesting Names
    oldnames LoopNesting
    loop) =
      Names -> LoopNesting -> Nesting
    Nesting (Names
    oldnames Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    newnames) LoopNesting
    loop
    -- ^ First pair element is the very innermost ("current") nest.  In
    -- the list, the outermost nest comes first.

    type Nestings = (Nesting, [Nesting])

    ppNestings :: Nestings -> String
    ppNestings :: Nestings -> [Char]
    ppNestings (Nesting
    nesting, [Nesting]
    nestings) =
      [[Char]] -> [Char]
    unlines ([[Char]] -> [Char]) -> [[Char]] -> [Char]
    forall a b. (a -> b) -> a -> b
    $ (Nesting -> [Char]) -> [Nesting] -> [[Char]]
    forall a b. (a -> b) -> [a] -> [b]
    map Nesting -> [Char]
    ppNesting ([Nesting] -> [[Char]]) -> [Nesting] -> [[Char]]
    forall a b. (a -> b) -> a -> b
    $ [Nesting]
    nestings [Nesting] -> [Nesting] -> [Nesting]
    forall a. [a] -> [a] -> [a]
    ++ [Nesting
    nesting]
      where
        ppNesting :: Nesting -> [Char]
    ppNesting (Nesting Names
    _ LoopNesting
    loop) = LoopNesting -> [Char]
    ppLoopNesting LoopNesting
    loop

    singleNesting :: Nesting -> Nestings
    singleNesting :: Nesting -> Nestings
    singleNesting = (,[])

    pushInnerNesting :: Nesting -> Nestings -> Nestings
    pushInnerNesting :: Nesting -> Nestings -> Nestings
    pushInnerNesting Nesting
    nesting (Nesting
    inner_nesting, [Nesting]
    nestings) =
      (Nesting
    nesting, [Nesting]
    nestings [Nesting] -> [Nesting] -> [Nesting]
    forall a. [a] -> [a] -> [a]
    ++ [Nesting
    inner_nesting])

    -- | Both parameters and let-bound.
    boundInNesting :: Nesting -> Names
    boundInNesting :: Nesting -> Names
    boundInNesting Nesting
    nesting =
      [VName] -> Names
    namesFromList ((Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName (LoopNesting -> [Param Type]
    loopNestingParams LoopNesting
    loop))
        Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Nesting -> Names
    nestingLetBound Nesting
    nesting
      where
        loop :: LoopNesting
    loop = Nesting -> LoopNesting
    nestingLoop Nesting
    nesting

    letBindInInnerNesting :: Names -> Nestings -> Nestings
    letBindInInnerNesting :: Names -> Nestings -> Nestings
    letBindInInnerNesting Names
    names (Nesting
    nest, [Nesting]
    nestings) =
      (Names -> Nesting -> Nesting
    letBindInNesting Names
    names Nesting
    nest, [Nesting]
    nestings)

    -- | Note: first element is *outermost* nesting.  This is different
    -- from the similar types elsewhere!
    type KernelNest = (LoopNesting, [LoopNesting])

    ppKernelNest :: KernelNest -> String
    ppKernelNest :: KernelNest -> [Char]
    ppKernelNest (LoopNesting
    nesting, [LoopNesting]
    nestings) =
      [[Char]] -> [Char]
    unlines ([[Char]] -> [Char]) -> [[Char]] -> [Char]
    forall a b. (a -> b) -> a -> b
    $ (LoopNesting -> [Char]) -> [LoopNesting] -> [[Char]]
    forall a b. (a -> b) -> [a] -> [b]
    map LoopNesting -> [Char]
    ppLoopNesting ([LoopNesting] -> [[Char]]) -> [LoopNesting] -> [[Char]]
    forall a b. (a -> b) -> a -> b
    $ LoopNesting
    nesting LoopNesting -> [LoopNesting] -> [LoopNesting]
    forall a. a -> [a] -> [a]
    : [LoopNesting]
    nestings

    -- | Retrieve the innermost kernel nesting.
    innermostKernelNesting :: KernelNest -> LoopNesting
    innermostKernelNesting :: KernelNest -> LoopNesting
    innermostKernelNesting (LoopNesting
    nest, [LoopNesting]
    nests) =
      LoopNesting -> Maybe LoopNesting -> LoopNesting
    forall a. a -> Maybe a -> a
    fromMaybe LoopNesting
    nest (Maybe LoopNesting -> LoopNesting)
    -> Maybe LoopNesting -> LoopNesting
    forall a b. (a -> b) -> a -> b
    $ [LoopNesting] -> Maybe LoopNesting
    forall a. [a] -> Maybe a
    maybeHead ([LoopNesting] -> Maybe LoopNesting)
    -> [LoopNesting] -> Maybe LoopNesting
    forall a b. (a -> b) -> a -> b
    $ [LoopNesting] -> [LoopNesting]
    forall a. [a] -> [a]
    reverse [LoopNesting]
    nests

    -- | Add new outermost nesting, pushing the current outermost to the
    -- list, also taking care to swap patterns if necessary.
    pushKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest
    pushKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest
    pushKernelNesting Target
    target LoopNesting
    newnest (LoopNesting
    nest, [LoopNesting]
    nests) =
      ( LoopNesting -> Target -> Pat Type -> LoopNesting
    fixNestingPatOrder LoopNesting
    newnest Target
    target (LoopNesting -> Pat Type
    loopNestingPat LoopNesting
    nest),
        LoopNesting
    nest LoopNesting -> [LoopNesting] -> [LoopNesting]
    forall a. a -> [a] -> [a]
    : [LoopNesting]
    nests
      )

    -- | Add new innermost nesting, pushing the current outermost to the
    -- list.  It is important that the 'Target' has the right order
    -- (non-permuted compared to what is expected by the outer nests).
    pushInnerKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest
    pushInnerKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest
    pushInnerKernelNesting Target
    target LoopNesting
    newnest (LoopNesting
    nest, [LoopNesting]
    nests) =
      (LoopNesting
    nest, [LoopNesting]
    nests [LoopNesting] -> [LoopNesting] -> [LoopNesting]
    forall a. [a] -> [a] -> [a]
    ++ [LoopNesting -> Target -> Pat Type -> LoopNesting
    fixNestingPatOrder LoopNesting
    newnest Target
    target (LoopNesting -> Pat Type
    loopNestingPat LoopNesting
    innermost)])
      where
        innermost :: LoopNesting
    innermost = case [LoopNesting] -> [LoopNesting]
    forall a. [a] -> [a]
    reverse [LoopNesting]
    nests of
          [] -> LoopNesting
    nest
          LoopNesting
    n : [LoopNesting]
    _ -> LoopNesting
    n

    fixNestingPatOrder :: LoopNesting -> Target -> Pat Type -> LoopNesting
    fixNestingPatOrder :: LoopNesting -> Target -> Pat Type -> LoopNesting
    fixNestingPatOrder LoopNesting
    nest (Pat Type
    _, Result
    res) Pat Type
    inner_pat =
      LoopNesting
    nest {loopNestingPat = basicPat pat'}
      where
        pat :: Pat Type
    pat = LoopNesting -> Pat Type
    loopNestingPat LoopNesting
    nest
        pat' :: [Ident]
    pat' = ((Ident, SubExpRes) -> Ident) -> [(Ident, SubExpRes)] -> [Ident]
    forall a b. (a -> b) -> [a] -> [b]
    map (Ident, SubExpRes) -> Ident
    forall a b. (a, b) -> a
    fst [(Ident, SubExpRes)]
    fixed_target
        fixed_target :: [(Ident, SubExpRes)]
    fixed_target = ((Ident, SubExpRes) -> Int)
    -> [(Ident, SubExpRes)] -> [(Ident, SubExpRes)]
    forall b a. Ord b => (a -> b) -> [a] -> [a]
    sortOn (Ident, SubExpRes) -> Int
    posInInnerPat ([(Ident, SubExpRes)] -> [(Ident, SubExpRes)])
    -> [(Ident, SubExpRes)] -> [(Ident, SubExpRes)]
    forall a b. (a -> b) -> a -> b
    $ [Ident] -> Result -> [(Ident, SubExpRes)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Pat Type -> [Ident]
    forall dec. Typed dec => Pat dec -> [Ident]
    patIdents Pat Type
    pat) Result
    res
        posInInnerPat :: (Ident, SubExpRes) -> Int
    posInInnerPat (Ident
    _, SubExpRes Certs
    _ (Var VName
    v)) = Int -> Maybe Int -> Int
    forall a. a -> Maybe a -> a
    fromMaybe Int
    0 (Maybe Int -> Int) -> Maybe Int -> Int
    forall a b. (a -> b) -> a -> b
    $ VName -> [VName] -> Maybe Int
    forall a. Eq a => a -> [a] -> Maybe Int
    elemIndex VName
    v ([VName] -> Maybe Int) -> [VName] -> Maybe Int
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    inner_pat
        posInInnerPat (Ident, SubExpRes)
    _ = Int
    0

    newKernel :: LoopNesting -> KernelNest
    newKernel :: LoopNesting -> KernelNest
    newKernel LoopNesting
    nest = (LoopNesting
    nest, [])

    kernelNestLoops :: KernelNest -> [LoopNesting]
    kernelNestLoops :: KernelNest -> [LoopNesting]
    kernelNestLoops (LoopNesting
    loop, [LoopNesting]
    loops) = LoopNesting
    loop LoopNesting -> [LoopNesting] -> [LoopNesting]
    forall a. a -> [a] -> [a]
    : [LoopNesting]
    loops

    scopeOfKernelNest :: (LParamInfo rep ~ Type) => KernelNest -> Scope rep
    scopeOfKernelNest :: forall rep. (LParamInfo rep ~ Type) => KernelNest -> Scope rep
    scopeOfKernelNest = (LoopNesting -> Scope rep) -> [LoopNesting] -> Scope rep
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap LoopNesting -> Scope rep
    forall rep. (LParamInfo rep ~ Type) => LoopNesting -> Scope rep
    scopeOfLoopNesting ([LoopNesting] -> Scope rep)
    -> (KernelNest -> [LoopNesting]) -> KernelNest -> Scope rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelNest -> [LoopNesting]
    kernelNestLoops

    boundInKernelNest :: KernelNest -> Names
    boundInKernelNest :: KernelNest -> Names
    boundInKernelNest = [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names)
    -> (KernelNest -> [Names]) -> KernelNest -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelNest -> [Names]
    boundInKernelNests

    boundInKernelNests :: KernelNest -> [Names]
    boundInKernelNests :: KernelNest -> [Names]
    boundInKernelNests =
      (LoopNesting -> Names) -> [LoopNesting] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map ([VName] -> Names
    namesFromList ([VName] -> Names)
    -> (LoopNesting -> [VName]) -> LoopNesting -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((Param Type, VName) -> VName) -> [(Param Type, VName)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param Type -> VName
    forall dec. Param dec -> VName
    paramName (Param Type -> VName)
    -> ((Param Type, VName) -> Param Type)
    -> (Param Type, VName)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Param Type, VName) -> Param Type
    forall a b. (a, b) -> a
    fst) ([(Param Type, VName)] -> [VName])
    -> (LoopNesting -> [(Param Type, VName)]) -> LoopNesting -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . LoopNesting -> [(Param Type, VName)]
    loopNestingParamsAndArrs)
        ([LoopNesting] -> [Names])
    -> (KernelNest -> [LoopNesting]) -> KernelNest -> [Names]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelNest -> [LoopNesting]
    kernelNestLoops

    kernelNestWidths :: KernelNest -> [SubExp]
    kernelNestWidths :: KernelNest -> [SubExp]
    kernelNestWidths = (LoopNesting -> SubExp) -> [LoopNesting] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map LoopNesting -> SubExp
    loopNestingWidth ([LoopNesting] -> [SubExp])
    -> (KernelNest -> [LoopNesting]) -> KernelNest -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelNest -> [LoopNesting]
    kernelNestLoops

    constructKernel ::
      (DistRep rep, MonadFreshNames m, LocalScope rep m) =>
      MkSegLevel rep m ->
      KernelNest ->
      Body rep ->
      m (Stm rep, Stms rep)
    constructKernel :: forall rep (m :: * -> *).
    (DistRep rep, MonadFreshNames m, LocalScope rep m) =>
    MkSegLevel rep m -> KernelNest -> Body rep -> m (Stm rep, Stms rep)
    constructKernel MkSegLevel rep m
    mk_lvl KernelNest
    kernel_nest Body rep
    inner_body = BuilderT rep m (Stm rep) -> m (Stm rep, Stms rep)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    BuilderT rep m a -> m (a, Stms rep)
    runBuilderT' (BuilderT rep m (Stm rep) -> m (Stm rep, Stms rep))
    -> BuilderT rep m (Stm rep) -> m (Stm rep, Stms rep)
    forall a b. (a -> b) -> a -> b
    $ do
      (ispace, inps) <- KernelNest -> BuilderT rep m ([(VName, SubExp)], [KernelInput])
    forall (m :: * -> *).
    MonadFreshNames m =>
    KernelNest -> m ([(VName, SubExp)], [KernelInput])
    flatKernel KernelNest
    kernel_nest
      let aux = LoopNesting -> StmAux ()
    loopNestingAux LoopNesting
    first_nest
          ispace_scope = [(VName, NameInfo rep)] -> Map VName (NameInfo rep)
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, NameInfo rep)] -> Map VName (NameInfo rep))
    -> [(VName, NameInfo rep)] -> Map VName (NameInfo rep)
    forall a b. (a -> b) -> a -> b
    $ ((VName, SubExp) -> (VName, NameInfo rep))
    -> [(VName, SubExp)] -> [(VName, NameInfo rep)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((,IntType -> NameInfo rep
    forall rep. IntType -> NameInfo rep
    IndexName IntType
    Int64) (VName -> (VName, NameInfo rep))
    -> ((VName, SubExp) -> VName)
    -> (VName, SubExp)
    -> (VName, NameInfo rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, SubExp) -> VName
    forall a b. (a, b) -> a
    fst) [(VName, SubExp)]
    ispace
          pat = LoopNesting -> Pat Type
    loopNestingPat LoopNesting
    first_nest
          rts = (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int -> Type -> Type
    forall u. Int -> TypeBase Shape u -> TypeBase Shape u
    stripArray ([(VName, SubExp)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(VName, SubExp)]
    ispace)) ([Type] -> [Type]) -> [Type] -> [Type]
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> [Type]
    forall dec. Typed dec => Pat dec -> [Type]
    patTypes Pat Type
    pat

      inner_body' <- fmap (uncurry (flip (KernelBody ()))) $
        runBuilder . localScope ispace_scope $ do
          mapM_ readKernelInput $ filter inputIsUsed inps
          res <- bodyBind inner_body
          forM res $ \(SubExpRes Certs
    cs SubExp
    se) -> KernelResult -> BuilderT rep (State VNameSource) KernelResult
    forall a. a -> BuilderT rep (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (KernelResult -> BuilderT rep (State VNameSource) KernelResult)
    -> KernelResult -> BuilderT rep (State VNameSource) KernelResult
    forall a b. (a -> b) -> a -> b
    $ ResultManifest -> Certs -> SubExp -> KernelResult
    Returns ResultManifest
    ResultMaySimplify Certs
    cs SubExp
    se

      (segop, aux_stms) <- lift $ mapKernel mk_lvl ispace [] rts inner_body'

      addStms aux_stms

      pure $ Let pat aux $ Op $ segOp segop
      where
        first_nest :: LoopNesting
    first_nest = KernelNest -> LoopNesting
    forall a b. (a, b) -> a
    fst KernelNest
    kernel_nest
        inputIsUsed :: KernelInput -> Bool
    inputIsUsed KernelInput
    input = KernelInput -> VName
    kernelInputName KernelInput
    input VName -> Names -> Bool
    `nameIn` Body rep -> Names
    forall a. FreeIn a => a -> Names
    freeIn Body rep
    inner_body

    -- | Flatten a kernel nesting to:
    --
    --  (1) The index space.
    --
    --  (2) The kernel inputs - note that some of these may be unused.
    flatKernel ::
      (MonadFreshNames m) =>
      KernelNest ->
      m ([(VName, SubExp)], [KernelInput])
    flatKernel :: forall (m :: * -> *).
    MonadFreshNames m =>
    KernelNest -> m ([(VName, SubExp)], [KernelInput])
    flatKernel (MapNesting Pat Type
    _ StmAux ()
    _ SubExp
    nesting_w [(Param Type, VName)]
    params_and_arrs, []) = do
      i <- [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "gtid"
      let inps =
            [ VName -> Type -> VName -> [SubExp] -> KernelInput
    KernelInput VName
    pname Type
    ptype VName
    arr [VName -> SubExp
    Var VName
    i]
              | (Param Attrs
    _ VName
    pname Type
    ptype, VName
    arr) <- [(Param Type, VName)]
    params_and_arrs
            ]
      pure ([(i, nesting_w)], inps)
    flatKernel (MapNesting Pat Type
    _ StmAux ()
    _ SubExp
    nesting_w [(Param Type, VName)]
    params_and_arrs, LoopNesting
    nest : [LoopNesting]
    nests) = do
      i <- [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "gtid"
      (ispace, inps) <- flatKernel (nest, nests)

      let inps' = (KernelInput -> KernelInput) -> [KernelInput] -> [KernelInput]
    forall a b. (a -> b) -> [a] -> [b]
    map KernelInput -> KernelInput
    fixupInput [KernelInput]
    inps
          isParam KernelInput
    inp =
            (Param Type, VName) -> VName
    forall a b. (a, b) -> b
    snd ((Param Type, VName) -> VName)
    -> Maybe (Param Type, VName) -> Maybe VName
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ((Param Type, VName) -> Bool)
    -> [(Param Type, VName)] -> Maybe (Param Type, VName)
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find ((VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == KernelInput -> VName
    kernelInputArray KernelInput
    inp) (VName -> Bool)
    -> ((Param Type, VName) -> VName) -> (Param Type, VName) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param Type -> VName
    forall dec. Param dec -> VName
    paramName (Param Type -> VName)
    -> ((Param Type, VName) -> Param Type)
    -> (Param Type, VName)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Param Type, VName) -> Param Type
    forall a b. (a, b) -> a
    fst) [(Param Type, VName)]
    params_and_arrs
          fixupInput KernelInput
    inp
            | Just VName
    arr <- KernelInput -> Maybe VName
    isParam KernelInput
    inp =
                KernelInput
    inp
                  { kernelInputArray = arr,
                    kernelInputIndices = Var i : kernelInputIndices inp
                  }
            | Bool
    otherwise =
                KernelInput
    inp

      pure ((i, nesting_w) : ispace, extra_inps i <> inps')
      where
        extra_inps :: VName -> [KernelInput]
    extra_inps VName
    i =
          [ VName -> Type -> VName -> [SubExp] -> KernelInput
    KernelInput VName
    pname Type
    ptype VName
    arr [VName -> SubExp
    Var VName
    i]
            | (Param Attrs
    _ VName
    pname Type
    ptype, VName
    arr) <- [(Param Type, VName)]
    params_and_arrs
          ]

    -- | Description of distribution to do.
    data DistributionBody = DistributionBody
      { DistributionBody -> Targets
    distributionTarget :: Targets,
        DistributionBody -> Names
    distributionFreeInBody :: Names,
        DistributionBody -> Map VName Ident
    distributionIdentityMap :: M.Map VName Ident,
        -- | Also related to avoiding identity mapping.
        DistributionBody -> Target -> Target
    distributionExpandTarget :: Target -> Target
      }

    distributionInnerPat :: DistributionBody -> Pat Type
    distributionInnerPat :: DistributionBody -> Pat Type
    distributionInnerPat = Target -> Pat Type
    forall a b. (a, b) -> a
    fst (Target -> Pat Type)
    -> (DistributionBody -> Target) -> DistributionBody -> Pat Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Targets -> Target
    innerTarget (Targets -> Target)
    -> (DistributionBody -> Targets) -> DistributionBody -> Target
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DistributionBody -> Targets
    distributionTarget

    distributionBodyFromStms ::
      (ASTRep rep) =>
      Targets ->
      Stms rep ->
      (DistributionBody, Result)
    distributionBodyFromStms :: forall rep.
    ASTRep rep =>
    Targets -> Stms rep -> (DistributionBody, Result)
    distributionBodyFromStms (Targets (Pat Type
    inner_pat, Result
    inner_res) [Target]
    targets) Stms rep
    stms =
      let bound_by_stms :: Names
    bound_by_stms = [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ Map VName (NameInfo rep) -> [VName]
    forall k a. Map k a -> [k]
    M.keys (Map VName (NameInfo rep) -> [VName])
    -> Map VName (NameInfo rep) -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Stms rep -> Map VName (NameInfo rep)
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Stms rep
    stms
          (Pat Type
    inner_pat', Result
    inner_res', Map VName Ident
    inner_identity_map, Target -> Target
    inner_expand_target) =
            Names
    -> Pat Type
    -> Result
    -> (Pat Type, Result, Map VName Ident, Target -> Target)
    removeIdentityMappingGeneral Names
    bound_by_stms Pat Type
    inner_pat Result
    inner_res
          free :: Names
    free =
            ((Stm rep -> Names) -> Stms rep -> Names
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Stm rep -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stms rep
    stms Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [Certs] -> Names
    forall a. FreeIn a => a -> Names
    freeIn ((SubExpRes -> Certs) -> Result -> [Certs]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> Certs
    resCerts Result
    inner_res))
              Names -> Names -> Names
    `namesSubtract` Names
    bound_by_stms
       in ( DistributionBody
              { distributionTarget :: Targets
    distributionTarget = Target -> [Target] -> Targets
    Targets (Pat Type
    inner_pat', Result
    inner_res') [Target]
    targets,
                distributionFreeInBody :: Names
    distributionFreeInBody = Names
    free,
                distributionIdentityMap :: Map VName Ident
    distributionIdentityMap = Map VName Ident
    inner_identity_map,
                distributionExpandTarget :: Target -> Target
    distributionExpandTarget = Target -> Target
    inner_expand_target
              },
            Result
    inner_res'
          )

    distributionBodyFromStm ::
      (ASTRep rep) =>
      Targets ->
      Stm rep ->
      (DistributionBody, Result)
    distributionBodyFromStm :: forall rep.
    ASTRep rep =>
    Targets -> Stm rep -> (DistributionBody, Result)
    distributionBodyFromStm Targets
    targets Stm rep
    stm =
      Targets -> Stms rep -> (DistributionBody, Result)
    forall rep.
    ASTRep rep =>
    Targets -> Stms rep -> (DistributionBody, Result)
    distributionBodyFromStms Targets
    targets (Stms rep -> (DistributionBody, Result))
    -> Stms rep -> (DistributionBody, Result)
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Stms rep
    forall rep. Stm rep -> Stms rep
    oneStm Stm rep
    stm

    createKernelNest ::
      forall rep m.
      (MonadFreshNames m, HasScope rep m) =>
      Nestings ->
      DistributionBody ->
      m (Maybe (Targets, KernelNest))
    createKernelNest :: forall rep (m :: * -> *).
    (MonadFreshNames m, HasScope rep m) =>
    Nestings -> DistributionBody -> m (Maybe (Targets, KernelNest))
    createKernelNest (Nesting
    inner_nest, [Nesting]
    nests) DistributionBody
    distrib_body = do
      let Targets Target
    target [Target]
    targets = DistributionBody -> Targets
    distributionTarget DistributionBody
    distrib_body
      Bool -> m () -> m ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Nesting] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Nesting]
    nests Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Target] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Target]
    targets) (m () -> m ()) -> m () -> m ()
    forall a b. (a -> b) -> a -> b
    $
        [Char] -> m ()
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> m ()) -> [Char] -> m ()
    forall a b. (a -> b) -> a -> b
    $
          [Char]
    "Nests and targets do not match!\n"
            [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "nests: "
            [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ Nestings -> [Char]
    ppNestings (Nesting
    inner_nest, [Nesting]
    nests)
            [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "\ntargets:"
            [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ Targets -> [Char]
    ppTargets (Target -> [Target] -> Targets
    Targets Target
    target [Target]
    targets)
      MaybeT m (Targets, KernelNest) -> m (Maybe (Targets, KernelNest))
    forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
    runMaybeT (MaybeT m (Targets, KernelNest) -> m (Maybe (Targets, KernelNest)))
    -> MaybeT m (Targets, KernelNest)
    -> m (Maybe (Targets, KernelNest))
    forall a b. (a -> b) -> a -> b
    $ ((KernelNest, Names, Targets) -> (Targets, KernelNest))
    -> MaybeT m (KernelNest, Names, Targets)
    -> MaybeT m (Targets, KernelNest)
    forall a b. (a -> b) -> MaybeT m a -> MaybeT m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (KernelNest, Names, Targets) -> (Targets, KernelNest)
    forall {b} {b} {a}. (b, b, a) -> (a, b)
    prepare (MaybeT m (KernelNest, Names, Targets)
     -> MaybeT m (Targets, KernelNest))
    -> MaybeT m (KernelNest, Names, Targets)
    -> MaybeT m (Targets, KernelNest)
    forall a b. (a -> b) -> a -> b
    $ [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets)
    recurse ([(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets))
    -> [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets)
    forall a b. (a -> b) -> a -> b
    $ [Nesting] -> [Target] -> [(Nesting, Target)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Nesting]
    nests [Target]
    targets
      where
        prepare :: (b, b, a) -> (a, b)
    prepare (b
    x, b
    _, a
    z) = (a
    z, b
    x)
        bound_in_nest :: Names
    bound_in_nest = [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names) -> [Names] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Nesting -> Names) -> [Nesting] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map Nesting -> Names
    boundInNesting ([Nesting] -> [Names]) -> [Nesting] -> [Names]
    forall a b. (a -> b) -> a -> b
    $ Nesting
    inner_nest Nesting -> [Nesting] -> [Nesting]
    forall a. a -> [a] -> [a]
    : [Nesting]
    nests
        distributableType :: Type -> Bool
    distributableType =
          (Names -> Names -> Bool
    forall a. Eq a => a -> a -> Bool
    == Names
    forall a. Monoid a => a
    mempty) (Names -> Bool) -> (Type -> Names) -> Type -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Names -> Names -> Names
    namesIntersection Names
    bound_in_nest (Names -> Names) -> (Type -> Names) -> Type -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [SubExp] -> Names
    forall a. FreeIn a => a -> Names
    freeIn ([SubExp] -> Names) -> (Type -> [SubExp]) -> Type -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Type -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims

        distributeAtNesting ::
          Nesting ->
          Pat Type ->
          (LoopNesting -> KernelNest, Names) ->
          M.Map VName Ident ->
          [Ident] ->
          (Target -> Targets) ->
          MaybeT m (KernelNest, Names, Targets)
        distributeAtNesting :: Nesting
    -> Pat Type
    -> (LoopNesting -> KernelNest, Names)
    -> Map VName Ident
    -> [Ident]
    -> (Target -> Targets)
    -> MaybeT m (KernelNest, Names, Targets)
    distributeAtNesting
          (Nesting Names
    nest_let_bound LoopNesting
    nest)
          Pat Type
    pat
          (LoopNesting -> KernelNest
    add_to_kernel, Names
    free_in_kernel)
          Map VName Ident
    identity_map
          [Ident]
    inner_returned_arrs
          Target -> Targets
    addTarget = do
            let nest' :: LoopNesting
    nest'@(MapNesting Pat Type
    _ StmAux ()
    aux SubExp
    w [(Param Type, VName)]
    params_and_arrs) =
                  Names -> LoopNesting -> LoopNesting
    removeUnusedNestingParts Names
    free_in_kernel LoopNesting
    nest
                ([Param Type]
    params, [VName]
    arrs) = [(Param Type, VName)] -> ([Param Type], [VName])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(Param Type, VName)]
    params_and_arrs
                param_names :: Names
    param_names = [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName [Param Type]
    params
                free_in_kernel' :: Names
    free_in_kernel' =
                  (LoopNesting -> Names
    forall a. FreeIn a => a -> Names
    freeIn LoopNesting
    nest' Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    free_in_kernel) Names -> Names -> Names
    `namesSubtract` Names
    param_names
                required_from_nest :: Names
    required_from_nest =
                  Names
    free_in_kernel' Names -> Names -> Names
    `namesIntersection` Names
    nest_let_bound

            required_from_nest_idents <-
              [VName] -> (VName -> MaybeT m Ident) -> MaybeT m [Ident]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM (Names -> [VName]
    namesToList Names
    required_from_nest) ((VName -> MaybeT m Ident) -> MaybeT m [Ident])
    -> (VName -> MaybeT m Ident) -> MaybeT m [Ident]
    forall a b. (a -> b) -> a -> b
    $ \VName
    name -> do
                t <- m Type -> MaybeT m Type
    forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m Type -> MaybeT m Type) -> m Type -> MaybeT m Type
    forall a b. (a -> b) -> a -> b
    $ VName -> m Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    name
                pure $ Ident name t

            (free_params, free_arrs, bind_in_target) <-
              fmap unzip3 $
                forM (inner_returned_arrs ++ required_from_nest_idents) $
                  \(Ident VName
    pname Type
    ptype) ->
                    case VName -> Map VName Ident -> Maybe Ident
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    pname Map VName Ident
    identity_map of
                      Maybe Ident
    Nothing -> do
                        arr <-
                          [Char] -> Type -> MaybeT m Ident
    forall (m :: * -> *).
    MonadFreshNames m =>
    [Char] -> Type -> m Ident
    newIdent (VName -> [Char]
    baseString VName
    pname [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_r") (Type -> MaybeT m Ident) -> Type -> MaybeT m Ident
    forall a b. (a -> b) -> a -> b
    $ Type -> SubExp -> Type
    forall d.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) NoUniqueness
    -> d -> TypeBase (ShapeBase d) NoUniqueness
    arrayOfRow Type
    ptype SubExp
    w
                        pure
                          ( Param mempty pname ptype,
                            arr,
                            True
                          )
                      Just Ident
    arr ->
                        (Param Type, Ident, Bool) -> MaybeT m (Param Type, Ident, Bool)
    forall a. a -> MaybeT m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                          ( Attrs -> VName -> Type -> Param Type
    forall dec. Attrs -> VName -> dec -> Param dec
    Param Attrs
    forall a. Monoid a => a
    mempty VName
    pname Type
    ptype,
                            Ident
    arr,
                            Bool
    False
                          )

            let free_arrs_pat =
                  [Ident] -> Pat Type
    basicPat ([Ident] -> Pat Type) -> [Ident] -> Pat Type
    forall a b. (a -> b) -> a -> b
    $ ((Bool, Ident) -> Ident) -> [(Bool, Ident)] -> [Ident]
    forall a b. (a -> b) -> [a] -> [b]
    map (Bool, Ident) -> Ident
    forall a b. (a, b) -> b
    snd ([(Bool, Ident)] -> [Ident]) -> [(Bool, Ident)] -> [Ident]
    forall a b. (a -> b) -> a -> b
    $ ((Bool, Ident) -> Bool) -> [(Bool, Ident)] -> [(Bool, Ident)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Bool, Ident) -> Bool
    forall a b. (a, b) -> a
    fst ([(Bool, Ident)] -> [(Bool, Ident)])
    -> [(Bool, Ident)] -> [(Bool, Ident)]
    forall a b. (a -> b) -> a -> b
    $ [Bool] -> [Ident] -> [(Bool, Ident)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Bool]
    bind_in_target [Ident]
    free_arrs
                free_params_pat =
                  ((Bool, Param Type) -> Param Type)
    -> [(Bool, Param Type)] -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Bool, Param Type) -> Param Type
    forall a b. (a, b) -> b
    snd ([(Bool, Param Type)] -> [Param Type])
    -> [(Bool, Param Type)] -> [Param Type]
    forall a b. (a -> b) -> a -> b
    $ ((Bool, Param Type) -> Bool)
    -> [(Bool, Param Type)] -> [(Bool, Param Type)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Bool, Param Type) -> Bool
    forall a b. (a, b) -> a
    fst ([(Bool, Param Type)] -> [(Bool, Param Type)])
    -> [(Bool, Param Type)] -> [(Bool, Param Type)]
    forall a b. (a -> b) -> a -> b
    $ [Bool] -> [Param Type] -> [(Bool, Param Type)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Bool]
    bind_in_target [Param Type]
    free_params

                (actual_params, actual_arrs) =
                  ( params ++ free_params,
                    arrs ++ map identName free_arrs
                  )
                actual_param_names =
                  [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName [Param Type]
    actual_params

                nest'' =
                  Names -> LoopNesting -> LoopNesting
    removeUnusedNestingParts Names
    free_in_kernel (LoopNesting -> LoopNesting) -> LoopNesting -> LoopNesting
    forall a b. (a -> b) -> a -> b
    $
                    Pat Type
    -> StmAux () -> SubExp -> [(Param Type, VName)] -> LoopNesting
    MapNesting Pat Type
    pat StmAux ()
    aux SubExp
    w ([(Param Type, VName)] -> LoopNesting)
    -> [(Param Type, VName)] -> LoopNesting
    forall a b. (a -> b) -> a -> b
    $
                      [Param Type] -> [VName] -> [(Param Type, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param Type]
    actual_params [VName]
    actual_arrs

                free_in_kernel'' =
                  (LoopNesting -> Names
    forall a. FreeIn a => a -> Names
    freeIn LoopNesting
    nest'' Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    free_in_kernel) Names -> Names -> Names
    `namesSubtract` Names
    actual_param_names

            unless
              ( all (distributableType . paramType) $
                  loopNestingParams nest''
              )
              $ fail "Would induce irregular array"
            pure
              ( add_to_kernel nest'',
                free_in_kernel'',
                addTarget (free_arrs_pat, varsRes $ map paramName free_params_pat)
              )

        recurse :: [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets)
        recurse :: [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets)
    recurse [] =
          Nesting
    -> Pat Type
    -> (LoopNesting -> KernelNest, Names)
    -> Map VName Ident
    -> [Ident]
    -> (Target -> Targets)
    -> MaybeT m (KernelNest, Names, Targets)
    distributeAtNesting
            Nesting
    inner_nest
            (DistributionBody -> Pat Type
    distributionInnerPat DistributionBody
    distrib_body)
            ( LoopNesting -> KernelNest
    newKernel,
              DistributionBody -> Names
    distributionFreeInBody DistributionBody
    distrib_body Names -> Names -> Names
    `namesIntersection` Names
    bound_in_nest
            )
            (DistributionBody -> Map VName Ident
    distributionIdentityMap DistributionBody
    distrib_body)
            []
            ((Target -> Targets) -> MaybeT m (KernelNest, Names, Targets))
    -> (Target -> Targets) -> MaybeT m (KernelNest, Names, Targets)
    forall a b. (a -> b) -> a -> b
    $ Target -> Targets
    singleTarget (Target -> Targets) -> (Target -> Target) -> Target -> Targets
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DistributionBody -> Target -> Target
    distributionExpandTarget DistributionBody
    distrib_body
        recurse ((Nesting
    nest, (Pat Type
    pat, Result
    res)) : [(Nesting, Target)]
    nests') = do
          (kernel@(outer, _), kernel_free, kernel_targets) <- [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets)
    recurse [(Nesting, Target)]
    nests'

          let (pat', res', identity_map, expand_target) =
                removeIdentityMappingFromNesting
                  (namesFromList $ patNames $ loopNestingPat outer)
                  pat
                  res

          distributeAtNesting
            nest
            pat'
            ( \LoopNesting
    k -> Target -> LoopNesting -> KernelNest -> KernelNest
    pushKernelNesting (Pat Type
    pat', Result
    res') LoopNesting
    k KernelNest
    kernel,
              kernel_free
            )
            identity_map
            (patIdents $ fst $ outerTarget kernel_targets)
            ((`pushOuterTarget` kernel_targets) . expand_target)

    removeUnusedNestingParts :: Names -> LoopNesting -> LoopNesting
    removeUnusedNestingParts :: Names -> LoopNesting -> LoopNesting
    removeUnusedNestingParts Names
    used (MapNesting Pat Type
    pat StmAux ()
    aux SubExp
    w [(Param Type, VName)]
    params_and_arrs) =
      Pat Type
    -> StmAux () -> SubExp -> [(Param Type, VName)] -> LoopNesting
    MapNesting Pat Type
    pat StmAux ()
    aux SubExp
    w ([(Param Type, VName)] -> LoopNesting)
    -> [(Param Type, VName)] -> LoopNesting
    forall a b. (a -> b) -> a -> b
    $ [Param Type] -> [VName] -> [(Param Type, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param Type]
    used_params [VName]
    used_arrs
      where
        ([Param Type]
    params, [VName]
    arrs) = [(Param Type, VName)] -> ([Param Type], [VName])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(Param Type, VName)]
    params_and_arrs
        ([Param Type]
    used_params, [VName]
    used_arrs) =
          [(Param Type, VName)] -> ([Param Type], [VName])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(Param Type, VName)] -> ([Param Type], [VName]))
    -> [(Param Type, VName)] -> ([Param Type], [VName])
    forall a b. (a -> b) -> a -> b
    $ ((Param Type, VName) -> Bool)
    -> [(Param Type, VName)] -> [(Param Type, VName)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter ((VName -> Names -> Bool
    `nameIn` Names
    used) (VName -> Bool)
    -> ((Param Type, VName) -> VName) -> (Param Type, VName) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param Type -> VName
    forall dec. Param dec -> VName
    paramName (Param Type -> VName)
    -> ((Param Type, VName) -> Param Type)
    -> (Param Type, VName)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Param Type, VName) -> Param Type
    forall a b. (a, b) -> a
    fst) ([(Param Type, VName)] -> [(Param Type, VName)])
    -> [(Param Type, VName)] -> [(Param Type, VName)]
    forall a b. (a -> b) -> a -> b
    $ [Param Type] -> [VName] -> [(Param Type, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param Type]
    params [VName]
    arrs

    removeIdentityMappingGeneral ::
      Names ->
      Pat Type ->
      Result ->
      ( Pat Type,
        Result,
        M.Map VName Ident,
        Target -> Target
      )
    removeIdentityMappingGeneral :: Names
    -> Pat Type
    -> Result
    -> (Pat Type, Result, Map VName Ident, Target -> Target)
    removeIdentityMappingGeneral Names
    bound Pat Type
    pat Result
    res =
      let ([(PatElem Type, (Certs, VName))]
    identities, [(PatElem Type, SubExpRes)]
    not_identities) =
            ((PatElem Type, SubExpRes)
     -> Either (PatElem Type, (Certs, VName)) (PatElem Type, SubExpRes))
    -> [(PatElem Type, SubExpRes)]
    -> ([(PatElem Type, (Certs, VName))], [(PatElem Type, SubExpRes)])
    forall a b c. (a -> Either b c) -> [a] -> ([b], [c])
    mapEither (PatElem Type, SubExpRes)
    -> Either (PatElem Type, (Certs, VName)) (PatElem Type, SubExpRes)
    isIdentity ([(PatElem Type, SubExpRes)]
     -> ([(PatElem Type, (Certs, VName))], [(PatElem Type, SubExpRes)]))
    -> [(PatElem Type, SubExpRes)]
    -> ([(PatElem Type, (Certs, VName))], [(PatElem Type, SubExpRes)])
    forall a b. (a -> b) -> a -> b
    $ [PatElem Type] -> Result -> [(PatElem Type, SubExpRes)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    pat) Result
    res
          ([PatElem Type]
    not_identity_patElems, Result
    not_identity_res) = [(PatElem Type, SubExpRes)] -> ([PatElem Type], Result)
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(PatElem Type, SubExpRes)]
    not_identities
          ([PatElem Type]
    identity_patElems, [(Certs, VName)]
    identity_res) = [(PatElem Type, (Certs, VName))]
    -> ([PatElem Type], [(Certs, VName)])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(PatElem Type, (Certs, VName))]
    identities
          expandTarget :: Target -> Target
    expandTarget (Pat Type
    tpat, Result
    tres) =
            ( [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
    $ Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    tpat [PatElem Type] -> [PatElem Type] -> [PatElem Type]
    forall a. [a] -> [a] -> [a]
    ++ [PatElem Type]
    identity_patElems,
              Result
    tres Result -> Result -> Result
    forall a. [a] -> [a] -> [a]
    ++ ((Certs, VName) -> SubExpRes) -> [(Certs, VName)] -> Result
    forall a b. (a -> b) -> [a] -> [b]
    map ((Certs -> SubExp -> SubExpRes) -> (Certs, SubExp) -> SubExpRes
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry Certs -> SubExp -> SubExpRes
    SubExpRes ((Certs, SubExp) -> SubExpRes)
    -> ((Certs, VName) -> (Certs, SubExp))
    -> (Certs, VName)
    -> SubExpRes
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName -> SubExp) -> (Certs, VName) -> (Certs, SubExp)
    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 -> SubExp
    Var) [(Certs, VName)]
    identity_res
            )
          identity_map :: Map VName Ident
    identity_map =
            [(VName, Ident)] -> Map VName Ident
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, Ident)] -> Map VName Ident)
    -> [(VName, Ident)] -> Map VName Ident
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [Ident] -> [(VName, Ident)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (((Certs, VName) -> VName) -> [(Certs, VName)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (Certs, VName) -> VName
    forall a b. (a, b) -> b
    snd [(Certs, VName)]
    identity_res) ([Ident] -> [(VName, Ident)]) -> [Ident] -> [(VName, Ident)]
    forall a b. (a -> b) -> a -> b
    $ (PatElem Type -> Ident) -> [PatElem Type] -> [Ident]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> Ident
    forall dec. Typed dec => PatElem dec -> Ident
    patElemIdent [PatElem Type]
    identity_patElems
       in ( [PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [PatElem Type]
    not_identity_patElems,
            Result
    not_identity_res,
            Map VName Ident
    identity_map,
            Target -> Target
    expandTarget
          )
      where
        isIdentity :: (PatElem Type, SubExpRes)
    -> Either (PatElem Type, (Certs, VName)) (PatElem Type, SubExpRes)
    isIdentity (PatElem Type
    patElem, SubExpRes Certs
    _ (Var VName
    v))
          | VName
    v VName -> Names -> Bool
    `notNameIn` Names
    bound = (PatElem Type, (Certs, VName))
    -> Either (PatElem Type, (Certs, VName)) (PatElem Type, SubExpRes)
    forall a b. a -> Either a b
    Left (PatElem Type
    patElem, (Certs
    forall a. Monoid a => a
    mempty, VName
    v))
        isIdentity (PatElem Type, SubExpRes)
    x = (PatElem Type, SubExpRes)
    -> Either (PatElem Type, (Certs, VName)) (PatElem Type, SubExpRes)
    forall a b. b -> Either a b
    Right (PatElem Type, SubExpRes)
    x

    removeIdentityMappingFromNesting ::
      Names ->
      Pat Type ->
      Result ->
      ( Pat Type,
        Result,
        M.Map VName Ident,
        Target -> Target
      )
    removeIdentityMappingFromNesting :: Names
    -> Pat Type
    -> Result
    -> (Pat Type, Result, Map VName Ident, Target -> Target)
    removeIdentityMappingFromNesting Names
    bound_in_nesting Pat Type
    pat Result
    res =
      let (Pat Type
    pat', Result
    res', Map VName Ident
    identity_map, Target -> Target
    expand_target) =
            Names
    -> Pat Type
    -> Result
    -> (Pat Type, Result, Map VName Ident, Target -> Target)
    removeIdentityMappingGeneral Names
    bound_in_nesting Pat Type
    pat Result
    res
       in (Pat Type
    pat', Result
    res', Map VName Ident
    identity_map, Target -> Target
    expand_target)

    tryDistribute ::
      ( DistRep rep,
        MonadFreshNames m,
        LocalScope rep m,
        MonadLogger m
      ) =>
      MkSegLevel rep m ->
      Nestings ->
      Targets ->
      Stms rep ->
      m (Maybe (Targets, Stms rep))
    tryDistribute :: forall rep (m :: * -> *).
    (DistRep rep, MonadFreshNames m, LocalScope rep m,
     MonadLogger m) =>
    MkSegLevel rep m
    -> Nestings -> Targets -> Stms rep -> m (Maybe (Targets, Stms rep))
    tryDistribute MkSegLevel rep m
    _ Nestings
    _ Targets
    targets Stms rep
    stms
      | Stms rep -> Bool
    forall a. Seq a -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null Stms rep
    stms =
          -- No point in distributing an empty kernel.
          Maybe (Targets, Stms rep) -> m (Maybe (Targets, Stms rep))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Maybe (Targets, Stms rep) -> m (Maybe (Targets, Stms rep)))
    -> Maybe (Targets, Stms rep) -> m (Maybe (Targets, Stms rep))
    forall a b. (a -> b) -> a -> b
    $ (Targets, Stms rep) -> Maybe (Targets, Stms rep)
    forall a. a -> Maybe a
    Just (Targets
    targets, Stms rep
    forall a. Monoid a => a
    mempty)
    tryDistribute MkSegLevel rep m
    mk_lvl Nestings
    nest Targets
    targets Stms rep
    stms =
      Nestings -> DistributionBody -> m (Maybe (Targets, KernelNest))
    forall rep (m :: * -> *).
    (MonadFreshNames m, HasScope rep m) =>
    Nestings -> DistributionBody -> m (Maybe (Targets, KernelNest))
    createKernelNest Nestings
    nest DistributionBody
    dist_body
        m (Maybe (Targets, KernelNest))
    -> (Maybe (Targets, KernelNest) -> m (Maybe (Targets, Stms rep)))
    -> m (Maybe (Targets, Stms rep))
    forall a b. m a -> (a -> m b) -> m b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \case
          Just (Targets
    targets', KernelNest
    distributed) -> do
            (kernel_stm, w_stms) <-
              Scope rep -> m (Stm rep, Stms rep) -> m (Stm rep, Stms rep)
    forall a. Scope rep -> m a -> m a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope (Targets -> Scope rep
    forall rep. DistRep rep => Targets -> Scope rep
    targetsScope Targets
    targets') (m (Stm rep, Stms rep) -> m (Stm rep, Stms rep))
    -> m (Stm rep, Stms rep) -> m (Stm rep, Stms rep)
    forall a b. (a -> b) -> a -> b
    $
                MkSegLevel rep m -> KernelNest -> Body rep -> m (Stm rep, Stms rep)
    forall rep (m :: * -> *).
    (DistRep rep, MonadFreshNames m, LocalScope rep m) =>
    MkSegLevel rep m -> KernelNest -> Body rep -> m (Stm rep, Stms rep)
    constructKernel MkSegLevel rep m
    mk_lvl KernelNest
    distributed (Body rep -> m (Stm rep, Stms rep))
    -> Body rep -> m (Stm rep, Stms rep)
    forall a b. (a -> b) -> a -> b
    $
                  Stms rep -> Result -> Body rep
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms rep
    stms Result
    inner_body_res
            distributed' <- renameStm kernel_stm
            logMsg $
              "distributing\n"
                ++ unlines (map prettyString $ stmsToList stms)
                ++ prettyString (snd $ innerTarget targets)
                ++ "\nas\n"
                ++ prettyString distributed'
                ++ "\ndue to targets\n"
                ++ ppTargets targets
                ++ "\nand with new targets\n"
                ++ ppTargets targets'
            pure $ Just (targets', w_stms <> oneStm distributed')
          Maybe (Targets, KernelNest)
    Nothing ->
            Maybe (Targets, Stms rep) -> m (Maybe (Targets, Stms rep))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Targets, Stms rep)
    forall a. Maybe a
    Nothing
      where
        (DistributionBody
    dist_body, Result
    inner_body_res) = Targets -> Stms rep -> (DistributionBody, Result)
    forall rep.
    ASTRep rep =>
    Targets -> Stms rep -> (DistributionBody, Result)
    distributionBodyFromStms Targets
    targets Stms rep
    stms

    tryDistributeStm ::
      (MonadFreshNames m, HasScope t m, ASTRep rep) =>
      Nestings ->
      Targets ->
      Stm rep ->
      m (Maybe (Result, Targets, KernelNest))
    tryDistributeStm :: forall (m :: * -> *) t rep.
    (MonadFreshNames m, HasScope t m, ASTRep rep) =>
    Nestings
    -> Targets -> Stm rep -> m (Maybe (Result, Targets, KernelNest))
    tryDistributeStm Nestings
    nest Targets
    targets Stm rep
    stm =
      ((Targets, KernelNest) -> (Result, Targets, KernelNest))
    -> Maybe (Targets, KernelNest)
    -> Maybe (Result, Targets, KernelNest)
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (Targets, KernelNest) -> (Result, Targets, KernelNest)
    addRes (Maybe (Targets, KernelNest)
     -> Maybe (Result, Targets, KernelNest))
    -> m (Maybe (Targets, KernelNest))
    -> m (Maybe (Result, Targets, KernelNest))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Nestings -> DistributionBody -> m (Maybe (Targets, KernelNest))
    forall rep (m :: * -> *).
    (MonadFreshNames m, HasScope rep m) =>
    Nestings -> DistributionBody -> m (Maybe (Targets, KernelNest))
    createKernelNest Nestings
    nest DistributionBody
    dist_body
      where
        (DistributionBody
    dist_body, Result
    res) = Targets -> Stm rep -> (DistributionBody, Result)
    forall rep.
    ASTRep rep =>
    Targets -> Stm rep -> (DistributionBody, Result)
    distributionBodyFromStm Targets
    targets Stm rep
    stm
        addRes :: (Targets, KernelNest) -> (Result, Targets, KernelNest)
    addRes (Targets
    targets', KernelNest
    kernel_nest) = (Result
    res, Targets
    targets', KernelNest
    kernel_nest)