SegMap.hs

-- | 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 is pe (Returns _ _ se) =
  copyDWIMFix (patElemName pe) (map Imp.le64 is) se []
writeResult _ _ res =
  error $ "writeResult: cannot handle " ++ prettyString res

compileSegMapBody ::
  Pat LetDecMem ->
  SegSpace ->
  KernelBody MCMem ->
  MulticoreGen ()
compileSegMapBody pat space (Body _ kstms kres) = do
  let (is, ns) = unzip $ unSegSpace space
      ns' = map pe64 ns
  dPrim_ (segFlat space) int64
  sOp $ Imp.GetTaskId (segFlat space)
  kstms' <- mapM renameStm kstms
  inISPC $
    generateChunkLoop "SegMap" Vectorized $ \i -> do
      dIndexSpace (zip is ns') i
      compileStms (freeIn kres) kstms' $
        zipWithM_ (writeResult is) (patElems pat) kres

compileSegMap ::
  Pat LetDecMem ->
  SegSpace ->
  KernelBody MCMem ->
  MulticoreGen ()
compileSegMap pat space kbody = do
  body <- collect $ compileSegMapBody pat space kbody
  free_params <- freeParams body
  emit $ Imp.Op $ Imp.ParLoop "segmap" body free_params