{-# LANGUAGE OverloadedStrings #-}

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

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 (Code Multicore) [Param]
  | -- | 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
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

-- | 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 (Code Multicore)

-- | 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 = Doc
"Dynamic"
  ppr Scheduling
Static = Doc
"Static"

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

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

instance Pretty Multicore where
  ppr :: Multicore -> Doc
ppr (GetLoopBounds VName
start VName
end) =
    (VName, VName) -> Doc
forall a. Pretty a => a -> Doc
ppr (VName
start, VName
end) Doc -> Doc -> Doc
<+> Doc
"<-" Doc -> Doc -> Doc
<+> Doc
"get_loop_bounds()"
  ppr (GetTaskId VName
v) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
<+> Doc
"<-" Doc -> Doc -> Doc
<+> Doc
"get_task_id()"
  ppr (GetNumTasks VName
v) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
<+> Doc
"<-" Doc -> Doc -> Doc
<+> Doc
"get_num_tasks()"
  ppr (SegOp String
s [Param]
free ParallelTask
seq_code Maybe ParallelTask
par_code [Param]
retval SchedulerInfo
scheduler) =
    Doc
"SegOp" Doc -> Doc -> Doc
<+> String -> Doc
text String
s Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" Doc
ppbody
    where
      ppbody :: Doc
ppbody =
        [Doc] -> Doc
stack
          [ SchedulerInfo -> Doc
forall a. Pretty a => a -> Doc
ppr SchedulerInfo
scheduler,
            String -> String -> Doc -> Doc
nestedBlock String
"free {" String
"}" ([Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
free),
            String -> String -> Doc -> Doc
nestedBlock String
"seq {" String
"}" (ParallelTask -> Doc
forall a. Pretty a => a -> Doc
ppr ParallelTask
seq_code),
            Doc -> (ParallelTask -> Doc) -> Maybe ParallelTask -> Doc
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Doc
forall a. Monoid a => a
mempty (String -> String -> Doc -> Doc
nestedBlock String
"par {" String
"}" (Doc -> Doc) -> (ParallelTask -> Doc) -> ParallelTask -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParallelTask -> Doc
forall a. Pretty a => a -> Doc
ppr) Maybe ParallelTask
par_code,
            String -> String -> Doc -> Doc
nestedBlock String
"retvals {" String
"}" ([Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
retval)
          ]
  ppr (ParLoop String
s Code Multicore
body [Param]
params) =
    Doc
"parloop" Doc -> Doc -> Doc
<+> String -> Doc
forall a. Pretty a => a -> Doc
ppr String
s Doc -> Doc -> Doc
</> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" Doc
ppbody
    where
      ppbody :: Doc
ppbody =
        [Doc] -> Doc
stack
          [ String -> String -> Doc -> Doc
nestedBlock String
"params {" String
"}" ([Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
params),
            String -> String -> Doc -> Doc
nestedBlock String
"body {" String
"}" (Code Multicore -> Doc
forall a. Pretty a => a -> Doc
ppr Code Multicore
body)
          ]
  ppr (Atomic AtomicOp
_) = Doc
"AtomicOp"

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 Code Multicore
code) = Code Multicore -> FV
forall a. FreeIn a => a -> FV
freeIn' Code Multicore
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
_ Code Multicore
body [Param]
_) =
    Code Multicore -> FV
forall a. FreeIn a => a -> FV
freeIn' Code Multicore
body
  freeIn' (Atomic AtomicOp
aop) = AtomicOp -> FV
forall a. FreeIn a => a -> FV
freeIn' AtomicOp
aop