-- | Multicore imperative code.
module Futhark.CodeGen.ImpCode.Multicore
  ( Program,
    Function,
    FunctionT (Function),
    Code,
    Multicore (..),
    Scheduling (..),
    SchedulerInfo (..),
    AtomicOp (..),
    ParallelTask (..),
    module Futhark.CodeGen.ImpCode,
  )
where

import Futhark.CodeGen.ImpCode hiding (Code, Function)
import qualified Futhark.CodeGen.ImpCode as Imp
import Futhark.Util.Pretty

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

-- | An imperative function.
type Function = Imp.Function Multicore

-- | A piece of imperative code, with multicore operations inside.
type Code = Imp.Code Multicore

-- | A multicore operation.
data Multicore
  = Segop String [Param] ParallelTask (Maybe ParallelTask) [Param] SchedulerInfo
  | ParLoop String VName Code Code Code [Param] VName
  | Atomic AtomicOp

-- | 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 (Imp.TExp Int32)) Exp
  | AtomicSub IntType VName VName (Count Elements (Imp.TExp Int32)) Exp
  | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int32)) Exp
  | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int32)) Exp
  | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int32)) Exp
  | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) Exp
  | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.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
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) = 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

data SchedulerInfo = SchedulerInfo
  { SchedulerInfo -> VName
nsubtasks :: VName, -- The variable that describes how many subtasks the scheduler created
    SchedulerInfo -> Exp
iterations :: Imp.Exp, -- The number of total iterations for a task
    SchedulerInfo -> Scheduling
scheduling :: Scheduling -- The type scheduling for the task
  }

data ParallelTask = ParallelTask
  { ParallelTask -> Code
task_code :: Code,
    ParallelTask -> VName
flatTid :: VName -- The variable for the thread id execution the code
  }

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

instance Pretty Scheduling where
  ppr :: Scheduling -> Doc
ppr Scheduling
Dynamic = String -> Doc
text String
"Dynamic"
  ppr Scheduling
Static = String -> Doc
text String
"Static"

-- TODO fix all of this!
instance Pretty SchedulerInfo where
  ppr :: SchedulerInfo -> Doc
ppr (SchedulerInfo VName
nsubtask Exp
i Scheduling
sched) =
    String -> Doc
text String
"SchedulingInfo"
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"number of subtasks"
      Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
nsubtask
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"scheduling"
      Doc -> Doc -> Doc
<+> Scheduling -> Doc
forall a. Pretty a => a -> Doc
ppr Scheduling
sched
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"iter"
      Doc -> Doc -> Doc
<+> Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
i

instance Pretty ParallelTask where
  ppr :: ParallelTask -> Doc
ppr (ParallelTask Code
code VName
_) =
    Code -> Doc
forall a. Pretty a => a -> Doc
ppr Code
code

instance Pretty Multicore where
  ppr :: Multicore -> Doc
ppr (Segop String
s [Param]
free ParallelTask
_par_code Maybe ParallelTask
seq_code [Param]
retval SchedulerInfo
scheduler) =
    String -> Doc
text String
"parfor"
      Doc -> Doc -> Doc
<+> SchedulerInfo -> Doc
forall a. Pretty a => a -> Doc
ppr SchedulerInfo
scheduler
      Doc -> Doc -> Doc
<+> [Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
free
      Doc -> Doc -> Doc
<+> String -> Doc
text String
s
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"seq_code"
      Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (Maybe ParallelTask -> Doc
forall a. Pretty a => a -> Doc
ppr Maybe ParallelTask
seq_code)
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"retvals"
      Doc -> Doc -> Doc
<+> [Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
retval
  ppr (ParLoop String
s VName
i Code
prebody Code
body Code
postbody [Param]
params VName
info) =
    String -> Doc
text String
"parloop" Doc -> Doc -> Doc
<+> String -> Doc
forall a. Pretty a => a -> Doc
ppr String
s Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
i
      Doc -> Doc -> Doc
<+> Code -> Doc
forall a. Pretty a => a -> Doc
ppr Code
prebody
      Doc -> Doc -> Doc
<+> [Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
params
      Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
info
      Doc -> Doc -> Doc
<+> Doc
langle
      Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (Code -> Doc
forall a. Pretty a => a -> Doc
ppr Code
body)
      Doc -> Doc -> Doc
<+> Code -> Doc
forall a. Pretty a => a -> Doc
ppr Code
postbody
  ppr (Atomic AtomicOp
_) = String -> Doc
text String
"AtomicOp"

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

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

instance FreeIn Multicore where
  freeIn' :: Multicore -> FV
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
_ VName
_ Code
prebody Code
body Code
postbody [Param]
_ VName
_) =
    Code -> FV
forall a. FreeIn a => a -> FV
freeIn' Code
prebody FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Names -> FV -> FV
fvBind (Code -> Names
forall a. Code a -> Names
Imp.declaredIn Code
prebody) (Code -> FV
forall a. FreeIn a => a -> FV
freeIn' (Code -> FV) -> Code -> FV
forall a b. (a -> b) -> a -> b
$ Code
body Code -> Code -> Code
forall a. Semigroup a => a -> a -> a
<> Code
postbody)
  freeIn' (Atomic AtomicOp
aop) = AtomicOp -> FV
forall a. FreeIn a => a -> FV
freeIn' AtomicOp
aop