-- | 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
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AtomicOp] -> ShowS
$cshowList :: [AtomicOp] -> ShowS
show :: AtomicOp -> String
$cshow :: AtomicOp -> String
showsPrec :: Int -> AtomicOp -> ShowS
$cshowsPrec :: Int -> AtomicOp -> ShowS
Show)

instance FreeIn AtomicOp where
  freeIn' :: AtomicOp -> FV
freeIn' (AtomicAdd IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
x
  freeIn' (AtomicSub IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
x
  freeIn' (AtomicAnd IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
x
  freeIn' (AtomicOr IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
x
  freeIn' (AtomicXor IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
x
  freeIn' (AtomicCmpXchg PrimType
_ VName
_ VName
arr Count Elements (TExp Int32)
i VName
retval Exp
x) = forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
x forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' VName
retval
  freeIn' (AtomicXchg PrimType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i forall a. Semigroup a => a -> a -> a
<> 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) =
    forall a. [Doc a] -> Doc a
stack
      [ forall a. Doc a -> Doc a -> Doc a -> Doc a
nestedBlock Doc ann
"scheduling {" Doc ann
"}" (forall a ann. Pretty a => a -> Doc ann
pretty Scheduling
sched),
        forall a. Doc a -> Doc a -> Doc a -> Doc a
nestedBlock Doc ann
"iter {" Doc ann
"}" (forall a ann. Pretty a => a -> Doc ann
pretty Exp
i)
      ]

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

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

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

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

instance FreeIn Multicore where
  freeIn' :: Multicore -> FV
freeIn' (GetLoopBounds VName
start VName
end) =
    forall a. FreeIn a => a -> FV
freeIn' (VName
start, VName
end)
  freeIn' (GetTaskId VName
v) =
    forall a. FreeIn a => a -> FV
freeIn' VName
v
  freeIn' (GetNumTasks VName
v) =
    forall a. FreeIn a => a -> FV
freeIn' VName
v
  freeIn' (SegOp String
_ [Param]
_ ParallelTask
par_code Maybe ParallelTask
seq_code [Param]
_ SchedulerInfo
info) =
    forall a. FreeIn a => a -> FV
freeIn' ParallelTask
par_code forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Maybe ParallelTask
seq_code forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' SchedulerInfo
info
  freeIn' (ParLoop String
_ MCCode
body [Param]
_) =
    forall a. FreeIn a => a -> FV
freeIn' MCCode
body
  freeIn' (Atomic AtomicOp
aop) =
    forall a. FreeIn a => a -> FV
freeIn' AtomicOp
aop
  freeIn' (ISPCKernel MCCode
body [Param]
_) =
    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) (forall a. FreeIn a => a -> FV
freeIn' MCCode
body forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
from forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
to)
  freeIn' (ForEachActive VName
i MCCode
body) =
    Names -> FV -> FV
fvBind (VName -> Names
oneName VName
i) (forall a. FreeIn a => a -> FV
freeIn' MCCode
body)
  freeIn' (ExtractLane VName
dest Exp
tar Exp
lane) =
    forall a. FreeIn a => a -> FV
freeIn' VName
dest forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp
tar forall a. Semigroup a => a -> a -> a
<> 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 =
  forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Names -> Bool
`notNameIn` Names
nonlexical)) forall a b. (a -> b) -> a -> b
$
    MCCode -> Map VName Space
declared forall a b. (a -> b) -> a -> b
$
      forall a. FunctionT a -> Code a
functionBody Function Multicore
func
  where
    nonlexical :: Names
nonlexical =
      MCCode -> Names
set (forall a. FunctionT a -> Code a
functionBody Function Multicore
func)
        forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList (forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName (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 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 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
_ = 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 -> forall a. Monoid a => a
mempty
    goOp MCCode -> a
_ Multicore
_ = forall a. Monoid a => a
mempty

    declared :: MCCode -> Map VName Space
declared (DeclareMem VName
mem Space
spc) =
      forall k a. k -> a -> Map k a
M.singleton VName
mem Space
spc
    declared MCCode
x = 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) = 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 {} = 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 -> forall a. Monoid a => a
mempty
        KernelHandling
OpaqueKernels -> [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
params forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
retvals
    set MCCode
x = forall {a}. Monoid a => (MCCode -> a) -> MCCode -> a
go MCCode -> Names
set MCCode
x