    -- | Multicore imperative code.
    module Futhark.CodeGen.ImpCode.Multicore
      ( Program,
        Multicore (..),
        MCCode,
        Scheduling (..),
        SchedulerInfo (..),
        AtomicOp (..),
        ParallelTask (..),
        KernelHandling (..),
        lexicalMemoryUsageMC,
        module Futhark.CodeGen.ImpCode,
      )
    where

    import Data.Map qualified as M
    import Futhark.CodeGen.ImpCode
    import Futhark.Util.Pretty

    -- | An imperative multicore program.
    type Program = Functions Multicore

    -- | A multicore operation.
    data Multicore
      = SegOp String [Param] ParallelTask (Maybe ParallelTask) [Param] SchedulerInfo
      | ParLoop String MCCode [Param]
      | -- | A kernel of ISPC code, or a scoped block in regular C.
        ISPCKernel MCCode [Param]
      | -- | A foreach loop in ISPC, or a regular for loop in C.
        ForEach VName Exp Exp MCCode
      | -- | A foreach_active loop in ISPC, or a single execution in C.
        ForEachActive VName MCCode
      | -- | Extract a value from a given lane and assign it to a variable.
        -- This is just a regular assignment in C.
        ExtractLane VName Exp Exp
      | -- | Retrieve inclusive start and exclusive end indexes of the
        -- chunk we are supposed to be executing.  Only valid immediately
        -- inside a 'ParLoop' construct!
        GetLoopBounds VName VName
      | -- | Retrieve the task ID that is currently executing.  Only valid
        -- immediately inside a 'ParLoop' construct!
        GetTaskId VName
      | -- | Retrieve the number of subtasks to execute.  Only valid
        -- immediately inside a 'SegOp' or 'ParLoop' construct!
        GetNumTasks VName
      | Atomic AtomicOp

    -- | Multicore code.
    type MCCode = Code Multicore

    -- | Atomic operations return the value stored before the update.
    -- This old value is stored in the first 'VName'.  The second 'VName'
    -- is the memory block to update.  The 'Exp' is the new value.
    data AtomicOp
      = AtomicAdd IntType VName VName (Count Elements (TExp Int32)) Exp
      | AtomicSub IntType VName VName (Count Elements (TExp Int32)) Exp
      | AtomicAnd IntType VName VName (Count Elements (TExp Int32)) Exp
      | AtomicOr IntType VName VName (Count Elements (TExp Int32)) Exp
      | AtomicXor IntType VName VName (Count Elements (TExp Int32)) Exp
      | AtomicXchg PrimType VName VName (Count Elements (TExp Int32)) Exp
      | AtomicCmpXchg PrimType VName VName (Count Elements (TExp Int32)) VName Exp
      deriving (Int -> AtomicOp -> ShowS
    [AtomicOp] -> ShowS
    AtomicOp -> String
    (Int -> AtomicOp -> ShowS)
    -> (AtomicOp -> String) -> ([AtomicOp] -> ShowS) -> Show AtomicOp
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> AtomicOp -> ShowS
    showsPrec :: Int -> AtomicOp -> ShowS
    $cshow :: AtomicOp -> String
    show :: AtomicOp -> String
    $cshowList :: [AtomicOp] -> ShowS
    showList :: [AtomicOp] -> ShowS
    Show)

    instance FreeIn AtomicOp where
      freeIn' :: AtomicOp -> FV
    freeIn' (AtomicAdd IntType
    _ VName
    _ VName
    arr Count Elements (TExp Int32)
    i Exp
    x) = VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    arr FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Count Elements (TExp Int32) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Count Elements (TExp Int32)
    i FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    x
      freeIn' (AtomicSub IntType
    _ VName
    _ VName
    arr Count Elements (TExp Int32)
    i Exp
    x) = VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    arr FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Count Elements (TExp Int32) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Count Elements (TExp Int32)
    i FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    x
      freeIn' (AtomicAnd IntType
    _ VName
    _ VName
    arr Count Elements (TExp Int32)
    i Exp
    x) = VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    arr FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Count Elements (TExp Int32) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Count Elements (TExp Int32)
    i FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    x
      freeIn' (AtomicOr IntType
    _ VName
    _ VName
    arr Count Elements (TExp Int32)
    i Exp
    x) = VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    arr FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Count Elements (TExp Int32) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Count Elements (TExp Int32)
    i FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    x
      freeIn' (AtomicXor IntType
    _ VName
    _ VName
    arr Count Elements (TExp Int32)
    i Exp
    x) = VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    arr FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Count Elements (TExp Int32) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Count Elements (TExp Int32)
    i FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    x
      freeIn' (AtomicCmpXchg PrimType
    _ VName
    _ VName
    arr Count Elements (TExp Int32)
    i VName
    retval Exp
    x) = VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    arr FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Count Elements (TExp Int32) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Count Elements (TExp Int32)
    i FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    x FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    retval
      freeIn' (AtomicXchg PrimType
    _ VName
    _ VName
    arr Count Elements (TExp Int32)
    i Exp
    x) = VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    arr FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Count Elements (TExp Int32) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Count Elements (TExp Int32)
    i FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    x

    -- | Information about parallel work that is do be done.  This is
    -- passed to the scheduler to help it make scheduling decisions.
    data SchedulerInfo = SchedulerInfo
      { -- | The number of total iterations for a task.
        SchedulerInfo -> Exp
    iterations :: Exp,
        -- | The type scheduling for the task.
        SchedulerInfo -> Scheduling
    scheduling :: Scheduling
      }

    -- | A task for a v'SegOp'.
    newtype ParallelTask = ParallelTask MCCode

    -- | Whether the Scheduler should schedule the tasks as Dynamic
    -- or it is restainted to Static
    data Scheduling
      = Dynamic
      | Static

    instance Pretty Scheduling where
      pretty :: forall ann. Scheduling -> Doc ann
    pretty Scheduling
    Dynamic = Doc ann
    "Dynamic"
      pretty Scheduling
    Static = Doc ann
    "Static"

    instance Pretty SchedulerInfo where
      pretty :: forall ann. SchedulerInfo -> Doc ann
    pretty (SchedulerInfo Exp
    i Scheduling
    sched) =
        [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    stack
          [ Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc ann
    "scheduling {" Doc ann
    "}" (Scheduling -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Scheduling -> Doc ann
    pretty Scheduling
    sched),
            Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc ann
    "iter {" Doc ann
    "}" (Exp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Exp -> Doc ann
    pretty Exp
    i)
          ]

    instance Pretty ParallelTask where
      pretty :: forall ann. ParallelTask -> Doc ann
    pretty (ParallelTask MCCode
    code) = MCCode -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. MCCode -> Doc ann
    pretty MCCode
    code

    instance Pretty Multicore where
      pretty :: forall ann. Multicore -> Doc ann
    pretty (GetLoopBounds VName
    start VName
    end) =
        (VName, VName) -> Doc ann
    forall ann. (VName, VName) -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty (VName
    start, VName
    end) Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "<-" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "get_loop_bounds()"
      pretty (GetTaskId VName
    v) =
        VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    v Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "<-" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "get_task_id()"
      pretty (GetNumTasks VName
    v) =
        VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    v Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "<-" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "get_num_tasks()"
      pretty (SegOp String
    s [Param]
    free ParallelTask
    seq_code Maybe ParallelTask
    par_code [Param]
    retval SchedulerInfo
    scheduler) =
        Doc ann
    "SegOp" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> String -> Doc ann
    forall ann. String -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty String
    s Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc ann
    "{" Doc ann
    "}" Doc ann
    forall {a}. Doc a
    ppbody
        where
          ppbody :: Doc a
    ppbody =
            [Doc a] -> Doc a
    forall a. [Doc a] -> Doc a
    stack
              [ SchedulerInfo -> Doc a
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SchedulerInfo -> Doc ann
    pretty SchedulerInfo
    scheduler,
                Doc a -> Doc a -> Doc a -> Doc a
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc a
    "free {" Doc a
    "}" ([Param] -> Doc a
    forall ann. [Param] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [Param]
    free),
                Doc a -> Doc a -> Doc a -> Doc a
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc a
    "seq {" Doc a
    "}" (ParallelTask -> Doc a
    forall a ann. Pretty a => a -> Doc ann
    forall ann. ParallelTask -> Doc ann
    pretty ParallelTask
    seq_code),
                Doc a -> (ParallelTask -> Doc a) -> Maybe ParallelTask -> Doc a
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Doc a
    forall a. Monoid a => a
    mempty (Doc a -> Doc a -> Doc a -> Doc a
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc a
    "par {" Doc a
    "}" (Doc a -> Doc a)
    -> (ParallelTask -> Doc a) -> ParallelTask -> Doc a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ParallelTask -> Doc a
    forall a ann. Pretty a => a -> Doc ann
    forall ann. ParallelTask -> Doc ann
    pretty) Maybe ParallelTask
    par_code,
                Doc a -> Doc a -> Doc a -> Doc a
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc a
    "retvals {" Doc a
    "}" ([Param] -> Doc a
    forall ann. [Param] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [Param]
    retval)
              ]
      pretty (ParLoop String
    s MCCode
    body [Param]
    params) =
        Doc ann
    "parloop" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> String -> Doc ann
    forall ann. String -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty String
    s Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc ann
    "{" Doc ann
    "}" Doc ann
    forall {a}. Doc a
    ppbody
        where
          ppbody :: Doc a
    ppbody =
            [Doc a] -> Doc a
    forall a. [Doc a] -> Doc a
    stack
              [ Doc a -> Doc a -> Doc a -> Doc a
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc a
    "params {" Doc a
    "}" ([Param] -> Doc a
    forall ann. [Param] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [Param]
    params),
                Doc a -> Doc a -> Doc a -> Doc a
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc a
    "body {" Doc a
    "}" (MCCode -> Doc a
    forall a ann. Pretty a => a -> Doc ann
    forall ann. MCCode -> Doc ann
    pretty MCCode
    body)
              ]
      pretty (Atomic AtomicOp
    _) =
        Doc ann
    "AtomicOp"
      pretty (ISPCKernel MCCode
    body [Param]
    _) =
        Doc ann
    "ispc" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc ann
    "{" Doc ann
    "}" (MCCode -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. MCCode -> Doc ann
    pretty MCCode
    body)
      pretty (ForEach VName
    i Exp
    from Exp
    to MCCode
    body) =
        Doc ann
    "foreach"
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    i
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "="
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Exp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Exp -> Doc ann
    pretty Exp
    from
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "to"
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Exp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Exp -> Doc ann
    pretty Exp
    to
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc ann
    "{" Doc ann
    "}" (MCCode -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. MCCode -> Doc ann
    pretty MCCode
    body)
      pretty (ForEachActive VName
    i MCCode
    body) =
        Doc ann
    "foreach_active"
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    i
          Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    nestedBlock Doc ann
    "{" Doc ann
    "}" (MCCode -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. MCCode -> Doc ann
    pretty MCCode
    body)
      pretty (ExtractLane VName
    dest Exp
    tar Exp
    lane) =
        VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    dest Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "<-" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "extract" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (Exp -> Doc ann) -> [Exp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map Exp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Exp -> Doc ann
    pretty [Exp
    tar, Exp
    lane])

    instance FreeIn SchedulerInfo where
      freeIn' :: SchedulerInfo -> FV
    freeIn' (SchedulerInfo Exp
    iter Scheduling
    _) = Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    iter

    instance FreeIn ParallelTask where
      freeIn' :: ParallelTask -> FV
    freeIn' (ParallelTask MCCode
    code) = MCCode -> FV
    forall a. FreeIn a => a -> FV
    freeIn' MCCode
    code

    instance FreeIn Multicore where
      freeIn' :: Multicore -> FV
    freeIn' (GetLoopBounds VName
    start VName
    end) =
        (VName, VName) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' (VName
    start, VName
    end)
      freeIn' (GetTaskId VName
    v) =
        VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    v
      freeIn' (GetNumTasks VName
    v) =
        VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    v
      freeIn' (SegOp String
    _ [Param]
    _ ParallelTask
    par_code Maybe ParallelTask
    seq_code [Param]
    _ SchedulerInfo
    info) =
        ParallelTask -> FV
    forall a. FreeIn a => a -> FV
    freeIn' ParallelTask
    par_code FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Maybe ParallelTask -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Maybe ParallelTask
    seq_code FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> SchedulerInfo -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SchedulerInfo
    info
      freeIn' (ParLoop String
    _ MCCode
    body [Param]
    _) =
        MCCode -> FV
    forall a. FreeIn a => a -> FV
    freeIn' MCCode
    body
      freeIn' (Atomic AtomicOp
    aop) =
        AtomicOp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' AtomicOp
    aop
      freeIn' (ISPCKernel MCCode
    body [Param]
    _) =
        MCCode -> FV
    forall a. FreeIn a => a -> FV
    freeIn' MCCode
    body
      freeIn' (ForEach VName
    i Exp
    from Exp
    to MCCode
    body) =
        Names -> FV -> FV
    fvBind (VName -> Names
    oneName VName
    i) (MCCode -> FV
    forall a. FreeIn a => a -> FV
    freeIn' MCCode
    body FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    from FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    to)
      freeIn' (ForEachActive VName
    i MCCode
    body) =
        Names -> FV -> FV
    fvBind (VName -> Names
    oneName VName
    i) (MCCode -> FV
    forall a. FreeIn a => a -> FV
    freeIn' MCCode
    body)
      freeIn' (ExtractLane VName
    dest Exp
    tar Exp
    lane) =
        VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn' VName
    dest FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    tar FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Exp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Exp
    lane

    -- | Whether 'lexicalMemoryUsageMC' should look inside nested kernels
    -- or not.
    data KernelHandling = TraverseKernels | OpaqueKernels

    -- | Like @lexicalMemoryUsage@, but traverses some inner multicore ops.
    lexicalMemoryUsageMC :: KernelHandling -> Function Multicore -> M.Map VName Space
    lexicalMemoryUsageMC :: KernelHandling -> Function Multicore -> Map VName Space
    lexicalMemoryUsageMC KernelHandling
    gokernel Function Multicore
    func =
      (VName -> Space -> Bool) -> Map VName Space -> Map VName Space
    forall k a. (k -> a -> Bool) -> Map k a -> Map k a
    M.filterWithKey (Bool -> Space -> Bool
    forall a b. a -> b -> a
    const (Bool -> Space -> Bool)
    -> (VName -> Bool) -> VName -> Space -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName -> Names -> Bool
    `notNameIn` Names
    nonlexical)) (Map VName Space -> Map VName Space)
    -> Map VName Space -> Map VName Space
    forall a b. (a -> b) -> a -> b
    $
        MCCode -> Map VName Space
    declared (MCCode -> Map VName Space) -> MCCode -> Map VName Space
    forall a b. (a -> b) -> a -> b
    $
          Function Multicore -> MCCode
    forall a. FunctionT a -> Code a
    functionBody Function Multicore
    func
      where
        nonlexical :: Names
    nonlexical =
          MCCode -> Names
    set (Function Multicore -> MCCode
    forall a. FunctionT a -> Code a
    functionBody Function Multicore
    func)
            Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [VName] -> Names
    namesFromList ((Param -> VName) -> [Param] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param -> VName
    paramName (Function Multicore -> [Param]
    forall a. FunctionT a -> [Param]
    functionOutput Function Multicore
    func))

        go :: (MCCode -> a) -> MCCode -> a
    go MCCode -> a
    f (MCCode
    x :>>: MCCode
    y) = MCCode -> a
    f MCCode
    x a -> a -> a
    forall a. Semigroup a => a -> a -> a
    <> MCCode -> a
    f MCCode
    y
        go MCCode -> a
    f (If TExp Bool
    _ MCCode
    x MCCode
    y) = MCCode -> a
    f MCCode
    x a -> a -> a
    forall a. Semigroup a => a -> a -> a
    <> MCCode -> a
    f MCCode
    y
        go MCCode -> a
    f (For VName
    _ Exp
    _ MCCode
    x) = MCCode -> a
    f MCCode
    x
        go MCCode -> a
    f (While TExp Bool
    _ MCCode
    x) = MCCode -> a
    f MCCode
    x
        go MCCode -> a
    f (Comment Text
    _ MCCode
    x) = MCCode -> a
    f MCCode
    x
        go MCCode -> a
    f (Op Multicore
    op) = (MCCode -> a) -> Multicore -> a
    goOp MCCode -> a
    f Multicore
    op
        go MCCode -> a
    _ MCCode
    _ = a
    forall a. Monoid a => a
    mempty

        -- We want SetMems and declarations to be visible through custom control flow
        -- so we don't erroneously treat a memblock that could be lexical as needing
        -- refcounting. Importantly, for ISPC, we do not look into kernels, since they
        -- go into new functions. For the Multicore backend, we can do it, though.
        goOp :: (MCCode -> a) -> Multicore -> a
    goOp MCCode -> a
    f (ForEach VName
    _ Exp
    _ Exp
    _ MCCode
    body) = (MCCode -> a) -> MCCode -> a
    go MCCode -> a
    f MCCode
    body
        goOp MCCode -> a
    f (ForEachActive VName
    _ MCCode
    body) = (MCCode -> a) -> MCCode -> a
    go MCCode -> a
    f MCCode
    body
        goOp MCCode -> a
    f (ISPCKernel MCCode
    body [Param]
    _) =
          case KernelHandling
    gokernel of
            KernelHandling
    TraverseKernels -> (MCCode -> a) -> MCCode -> a
    go MCCode -> a
    f MCCode
    body
            KernelHandling
    OpaqueKernels -> a
    forall a. Monoid a => a
    mempty
        goOp MCCode -> a
    _ Multicore
    _ = a
    forall a. Monoid a => a
    mempty

        declared :: MCCode -> Map VName Space
    declared (DeclareMem VName
    mem Space
    spc) =
          VName -> Space -> Map VName Space
    forall k a. k -> a -> Map k a
    M.singleton VName
    mem Space
    spc
        declared MCCode
    x = (MCCode -> Map VName Space) -> MCCode -> Map VName Space
    forall {a}. Monoid a => (MCCode -> a) -> MCCode -> a
    go MCCode -> Map VName Space
    declared MCCode
    x

        set :: MCCode -> Names
    set (SetMem VName
    x VName
    y Space
    _) = [VName] -> Names
    namesFromList [VName
    x, VName
    y]
        set (Call [VName]
    _ Name
    _ [Arg]
    args) = (Arg -> Names) -> [Arg] -> Names
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Arg -> Names
    onArg [Arg]
    args
          where
            onArg :: Arg -> Names
    onArg ExpArg {} = Names
    forall a. Monoid a => a
    mempty
            onArg (MemArg VName
    x) = VName -> Names
    oneName VName
    x
        -- Critically, don't treat inputs to nested segops as lexical when generating
        -- ISPC, since we want to use AoS memory for lexical blocks, which is
        -- incompatible with pointer assignmentes visible in C.
        set (Op (SegOp String
    _ [Param]
    params ParallelTask
    _ Maybe ParallelTask
    _ [Param]
    retvals SchedulerInfo
    _)) =
          case KernelHandling
    gokernel of
            KernelHandling
    TraverseKernels -> Names
    forall a. Monoid a => a
    mempty
            KernelHandling
    OpaqueKernels -> [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Param -> VName) -> [Param] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param -> VName
    paramName [Param]
    params [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> (Param -> VName) -> [Param] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param -> VName
    paramName [Param]
    retvals
        set MCCode
    x = (MCCode -> Names) -> MCCode -> Names
    forall {a}. Monoid a => (MCCode -> a) -> MCCode -> a
    go MCCode -> Names
    set MCCode
    x