-- | Multicore code generation for 'SegMap'.
module Futhark.CodeGen.ImpGen.Multicore.SegMap
  ( compileSegMap,
  )
where

import Control.Monad
import Futhark.CodeGen.ImpCode.Multicore qualified as Imp
import Futhark.CodeGen.ImpGen
import Futhark.CodeGen.ImpGen.Multicore.Base
import Futhark.IR.MCMem
import Futhark.Transform.Rename

writeResult ::
  [VName] ->
  PatElem dec ->
  KernelResult ->
  MulticoreGen ()
writeResult :: forall dec.
[VName] -> PatElem dec -> KernelResult -> MulticoreGen ()
writeResult [VName]
is PatElem dec
pe (Returns ResultManifest
_ Certs
_ SubExp
se) =
  VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (PatElem dec -> VName
forall dec. PatElem dec -> VName
patElemName PatElem dec
pe) ((VName -> TExp Int64) -> [VName] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TExp Int64
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
is) SubExp
se []
writeResult [VName]
_ PatElem dec
_ KernelResult
res =
  [Char] -> MulticoreGen ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> MulticoreGen ()) -> [Char] -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ [Char]
"writeResult: cannot handle " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ KernelResult -> [Char]
forall a. Pretty a => a -> [Char]
prettyString KernelResult
res

compileSegMapBody ::
  Pat LetDecMem ->
  SegSpace ->
  KernelBody MCMem ->
  MulticoreGen ()
compileSegMapBody :: Pat LetDecMem -> SegSpace -> KernelBody MCMem -> MulticoreGen ()
compileSegMapBody Pat LetDecMem
pat SegSpace
space (Body BodyDec MCMem
_ Stms MCMem
kstms [KernelResult]
kres) = do
  let ([VName]
is, [SubExp]
ns) = [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, SubExp)] -> ([VName], [SubExp]))
-> [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
      ns' :: [TExp Int64]
ns' = (SubExp -> TExp Int64) -> [SubExp] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TExp Int64
pe64 [SubExp]
ns
  VName -> PrimType -> MulticoreGen ()
forall rep r op. VName -> PrimType -> ImpM rep r op ()
dPrim_ (SegSpace -> VName
segFlat SegSpace
space) PrimType
int64
  Multicore -> MulticoreGen ()
forall op rep r. op -> ImpM rep r op ()
sOp (Multicore -> MulticoreGen ()) -> Multicore -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ VName -> Multicore
Imp.GetTaskId (SegSpace -> VName
segFlat SegSpace
space)
  kstms' <- (Stm MCMem -> ImpM MCMem HostEnv Multicore (Stm MCMem))
-> Stms MCMem -> ImpM MCMem HostEnv Multicore (Stms MCMem)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Seq a -> m (Seq b)
mapM Stm MCMem -> ImpM MCMem HostEnv Multicore (Stm MCMem)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Stm rep -> m (Stm rep)
renameStm Stms MCMem
kstms
  inISPC $
    generateChunkLoop "SegMap" Vectorized $ \TExp Int64
i -> do
      [(VName, TExp Int64)] -> TExp Int64 -> MulticoreGen ()
forall rep r op.
[(VName, TExp Int64)] -> TExp Int64 -> ImpM rep r op ()
dIndexSpace ([VName] -> [TExp Int64] -> [(VName, TExp Int64)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
is [TExp Int64]
ns') TExp Int64
i
      Names -> Stms MCMem -> MulticoreGen () -> MulticoreGen ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms ([KernelResult] -> Names
forall a. FreeIn a => a -> Names
freeIn [KernelResult]
kres) Stms MCMem
kstms' (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$
        (PatElem LetDecMem -> KernelResult -> MulticoreGen ())
-> [PatElem LetDecMem] -> [KernelResult] -> MulticoreGen ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ ([VName] -> PatElem LetDecMem -> KernelResult -> MulticoreGen ()
forall dec.
[VName] -> PatElem dec -> KernelResult -> MulticoreGen ()
writeResult [VName]
is) (Pat LetDecMem -> [PatElem LetDecMem]
forall dec. Pat dec -> [PatElem dec]
patElems Pat LetDecMem
pat) [KernelResult]
kres

compileSegMap ::
  Pat LetDecMem ->
  SegSpace ->
  KernelBody MCMem ->
  MulticoreGen ()
compileSegMap :: Pat LetDecMem -> SegSpace -> KernelBody MCMem -> MulticoreGen ()
compileSegMap Pat LetDecMem
pat SegSpace
space KernelBody MCMem
kbody = do
  body <- MulticoreGen () -> ImpM MCMem HostEnv Multicore (Code Multicore)
forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
collect (MulticoreGen () -> ImpM MCMem HostEnv Multicore (Code Multicore))
-> MulticoreGen () -> ImpM MCMem HostEnv Multicore (Code Multicore)
forall a b. (a -> b) -> a -> b
$ Pat LetDecMem -> SegSpace -> KernelBody MCMem -> MulticoreGen ()
compileSegMapBody Pat LetDecMem
pat SegSpace
space KernelBody MCMem
kbody
  free_params <- freeParams body
  emit $ Imp.Op $ Imp.ParLoop "segmap" body free_params