{-# LANGUAGE QuasiQuotes #-}

-- | C code generator.  This module can convert a correct ImpCode
-- program to an equivalent ISPC program.
module Futhark.CodeGen.Backends.MulticoreISPC
  ( compileProg,
    GC.CParts (..),
    GC.asLibrary,
    GC.asExecutable,
    GC.asServer,
    operations,
    ISPCState,
  )
where

import Control.Lens (each, over)
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.DList qualified as DL
import Data.List (unzip4)
import Data.Map qualified as M
import Data.Maybe
import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericC qualified as GC
import Futhark.CodeGen.Backends.GenericC.Pretty
import Futhark.CodeGen.Backends.MulticoreC qualified as MC
import Futhark.CodeGen.Backends.MulticoreC.Boilerplate (generateBoilerplate)
import Futhark.CodeGen.Backends.SimpleRep
import Futhark.CodeGen.ImpCode.Multicore
import Futhark.CodeGen.ImpGen.Multicore qualified as ImpGen
import Futhark.CodeGen.RTS.C (errorsH, ispcUtilH, uniformH)
import Futhark.IR.MCMem (MCMem, Prog)
import Futhark.IR.Prop (isBuiltInFunction)
import Futhark.MonadFreshNames
import Language.C.Quote.OpenCL qualified as C
import Language.C.Syntax qualified as C
import NeatInterpolation (untrimming)

type ISPCCompilerM a = GC.CompilerM Multicore ISPCState a

-- | Transient state tracked by the ISPC backend.
data ISPCState = ISPCState
  { ISPCState -> DList Definition
sDefs :: DL.DList C.Definition,
    ISPCState -> Names
sUniform :: Names
  }

uniform :: C.TypeQual
uniform :: TypeQual
uniform = String -> SrcLoc -> TypeQual
C.EscTypeQual String
"uniform" SrcLoc
forall a. IsLocation a => a
noLoc

unmasked :: C.TypeQual
unmasked :: TypeQual
unmasked = String -> SrcLoc -> TypeQual
C.EscTypeQual String
"unmasked" SrcLoc
forall a. IsLocation a => a
noLoc

export :: C.TypeQual
export :: TypeQual
export = String -> SrcLoc -> TypeQual
C.EscTypeQual String
"export" SrcLoc
forall a. IsLocation a => a
noLoc

varying :: C.TypeQual
varying :: TypeQual
varying = String -> SrcLoc -> TypeQual
C.EscTypeQual String
"varying" SrcLoc
forall a. IsLocation a => a
noLoc

-- | Compile the program to C and ISPC code using multicore operations.
compileProg ::
  (MonadFreshNames m) => T.Text -> Prog MCMem -> m (ImpGen.Warnings, (GC.CParts, T.Text))
compileProg :: forall (m :: * -> *).
MonadFreshNames m =>
Text -> Prog MCMem -> m (Warnings, (CParts, Text))
compileProg Text
version Prog MCMem
prog = do
  -- Dynamic scheduling seems completely broken currently, so we disable it.
  (ws, defs) <- Prog MCMem -> m (Warnings, Definitions Multicore)
forall (m :: * -> *).
MonadFreshNames m =>
Prog MCMem -> m (Warnings, Definitions Multicore)
ImpGen.compileProg Prog MCMem
prog
  let Functions funs = defFuns defs

  (ws', (cparts, endstate)) <-
    traverse
      ( GC.compileProg'
          "ispc"
          version
          mempty
          operations
          (ISPCState mempty mempty)
          ( do
              generateBoilerplate
              mapM_ compileBuiltinFun funs
          )
          "#include <pthread.h>\n"
          (DefaultSpace, [DefaultSpace])
          MC.cliOptions
      )
      (ws, defs)

  let ispc_decls = [Definition] -> Text
definitionsText ([Definition] -> Text) -> [Definition] -> Text
forall a b. (a -> b) -> a -> b
$ DList Definition -> [Definition]
forall a. DList a -> [a]
DL.toList (DList Definition -> [Definition])
-> DList Definition -> [Definition]
forall a b. (a -> b) -> a -> b
$ ISPCState -> DList Definition
sDefs (ISPCState -> DList Definition) -> ISPCState -> DList Definition
forall a b. (a -> b) -> a -> b
$ CompilerState ISPCState -> ISPCState
forall s. CompilerState s -> s
GC.compUserState CompilerState ISPCState
endstate

  -- The bool #define is a workaround around an ISPC bug, stdbool doesn't get included.
  let ispcdefs =
        [untrimming|
#define bool uint8
typedef int64 int64_t;
typedef int32 int32_t;
typedef int16 int16_t;
typedef int8 int8_t;
typedef int8 char;
typedef unsigned int64 uint64_t;
typedef unsigned int32 uint32_t;
typedef unsigned int16 uint16_t;
typedef unsigned int8 uint8_t;
#define volatile
#define SCALAR_FUN_ATTR static inline

$errorsH

#define INFINITY (floatbits((uniform int)0x7f800000))
#define NAN (floatbits((uniform int)0x7fc00000))
#define fabs(x) abs(x)
#define FUTHARK_F64_ENABLED
$cScalarDefs

$uniformH

$ispcUtilH

$ispc_decls|]

  pure (ws', (cparts, ispcdefs))

-- | Compiler operations specific to the ISPC multicore backend.
operations :: GC.Operations Multicore ISPCState
operations :: Operations Multicore ISPCState
operations =
  Operations Multicore ISPCState
forall s. Operations Multicore s
MC.operations
    { GC.opsCompiler = compileOp,
      -- FIXME: the default codegen for LMAD copies does not work for ISPC.
      GC.opsCopies = mempty
    }

ispcDecl :: C.Definition -> ISPCCompilerM ()
ispcDecl :: Definition -> CompilerM Multicore ISPCState ()
ispcDecl Definition
def =
  (ISPCState -> ISPCState) -> CompilerM Multicore ISPCState ()
forall s op. (s -> s) -> CompilerM op s ()
GC.modifyUserState (\ISPCState
s -> ISPCState
s {sDefs = sDefs s <> DL.singleton def})

ispcEarlyDecl :: C.Definition -> ISPCCompilerM ()
ispcEarlyDecl :: Definition -> CompilerM Multicore ISPCState ()
ispcEarlyDecl Definition
def =
  (ISPCState -> ISPCState) -> CompilerM Multicore ISPCState ()
forall s op. (s -> s) -> CompilerM op s ()
GC.modifyUserState (\ISPCState
s -> ISPCState
s {sDefs = DL.singleton def <> sDefs s})

ispcDef :: MC.DefSpecifier ISPCState
ispcDef :: DefSpecifier ISPCState
ispcDef String
s Name -> CompilerM Multicore ISPCState Definition
f = do
  s' <- String -> CompilerM Multicore ISPCState Name
forall op s. String -> CompilerM op s Name
MC.multicoreName String
s
  ispcDecl =<< f s'
  pure s'

-- | Expose a struct to both ISPC and C.
sharedDef :: MC.DefSpecifier ISPCState
sharedDef :: DefSpecifier ISPCState
sharedDef String
s Name -> CompilerM Multicore ISPCState Definition
f = do
  s' <- String -> CompilerM Multicore ISPCState Name
forall op s. String -> CompilerM op s Name
MC.multicoreName String
s
  ispcDecl =<< f s'
  GC.earlyDecl =<< f s'
  pure s'

-- | ISPC has no string literals, so this makes one in C and exposes it via an
-- external function, returning the name.
makeStringLiteral :: String -> ISPCCompilerM Name
makeStringLiteral :: String -> CompilerM Multicore ISPCState Name
makeStringLiteral String
str = do
  name <- DefSpecifier ISPCState
forall s. DefSpecifier s
MC.multicoreDef String
"strlit_shim" ((Name -> CompilerM Multicore ISPCState Definition)
 -> CompilerM Multicore ISPCState Name)
-> (Name -> CompilerM Multicore ISPCState Definition)
-> CompilerM Multicore ISPCState Name
forall a b. (a -> b) -> a -> b
$ \Name
s ->
    Definition -> CompilerM Multicore ISPCState Definition
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cedecl|char* $id:s() { return $string:str; }|]
  ispcDecl
    [C.cedecl|extern "C" $tyqual:unmasked $tyqual:uniform char* $tyqual:uniform $id:name();|]
  pure name

-- | Set memory in ISPC
setMem :: (C.ToExp a, C.ToExp b) => a -> b -> Space -> ISPCCompilerM ()
setMem :: forall a b.
(ToExp a, ToExp b) =>
a -> b -> Space -> CompilerM Multicore ISPCState ()
setMem a
dest b
src Space
space = do
  let src_s :: String
src_s = Text -> String
T.unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ Exp -> Text
expText (Exp -> Text) -> Exp -> Text
forall a b. (a -> b) -> a -> b
$ b -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
C.toExp b
src SrcLoc
forall a. IsLocation a => a
noLoc
  strlit <- String -> CompilerM Multicore ISPCState Name
makeStringLiteral String
src_s
  GC.stm
    [C.cstm|if ($id:(GC.fatMemSet space)(ctx, &$exp:dest, &$exp:src,
                                            $id:strlit()) != 0) {
                    $escstm:("unmasked { return 1; }")
                  }|]

-- | Unref memory in ISPC
unRefMem :: (C.ToExp a) => a -> Space -> ISPCCompilerM ()
unRefMem :: forall a. ToExp a => a -> Space -> CompilerM Multicore ISPCState ()
unRefMem a
mem Space
space = do
  cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM Multicore ISPCState (Maybe VName)
-> CompilerM Multicore ISPCState Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> CompilerM Multicore ISPCState (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
GC.cacheMem a
mem
  let mem_s = Text -> String
T.unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ Exp -> Text
expText (Exp -> Text) -> Exp -> Text
forall a b. (a -> b) -> a -> b
$ a -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
C.toExp a
mem SrcLoc
forall a. IsLocation a => a
noLoc
  strlit <- makeStringLiteral mem_s
  unless cached $
    GC.stm
      [C.cstm|if ($id:(GC.fatMemUnRef space)(ctx, &$exp:mem, $id:strlit()) != 0) {
                  $escstm:("unmasked { return 1; }")
                }|]

-- | Allocate memory in ISPC
allocMem ::
  (C.ToExp a, C.ToExp b) =>
  a ->
  b ->
  Space ->
  C.Stm ->
  ISPCCompilerM ()
allocMem :: forall a b.
(ToExp a, ToExp b) =>
a -> b -> Space -> Stm -> CompilerM Multicore ISPCState ()
allocMem a
mem b
size Space
space Stm
on_failure = do
  let mem_s :: String
mem_s = Text -> String
T.unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ Exp -> Text
expText (Exp -> Text) -> Exp -> Text
forall a b. (a -> b) -> a -> b
$ a -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
C.toExp a
mem SrcLoc
forall a. IsLocation a => a
noLoc
  strlit <- String -> CompilerM Multicore ISPCState Name
makeStringLiteral String
mem_s
  GC.stm
    [C.cstm|if ($id:(GC.fatMemAlloc space)(ctx, &$exp:mem, $exp:size, $id:strlit())) {
                    $stm:on_failure
            }|]

-- | Free memory in ISPC
freeAllocatedMem :: ISPCCompilerM [C.BlockItem]
freeAllocatedMem :: ISPCCompilerM [BlockItem]
freeAllocatedMem = CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
GC.collect (CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem])
-> CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$ ((VName, Space) -> CompilerM Multicore ISPCState ())
-> [(VName, Space)] -> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((VName -> Space -> CompilerM Multicore ISPCState ())
-> (VName, Space) -> CompilerM Multicore ISPCState ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry VName -> Space -> CompilerM Multicore ISPCState ()
forall a. ToExp a => a -> Space -> CompilerM Multicore ISPCState ()
unRefMem) ([(VName, Space)] -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState [(VName, Space)]
-> CompilerM Multicore ISPCState ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (CompilerState ISPCState -> [(VName, Space)])
-> CompilerM Multicore ISPCState [(VName, Space)]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState ISPCState -> [(VName, Space)]
forall s. CompilerState s -> [(VName, Space)]
GC.compDeclaredMem

-- | Given a ImpCode function, generate all the required machinery for calling
-- it in ISPC, both in a varying or uniform context. This involves handling
-- for the fact that ISPC cannot pass structs by value to external functions.
compileBuiltinFun :: (Name, Function op) -> ISPCCompilerM ()
compileBuiltinFun :: forall op. (Name, Function op) -> CompilerM Multicore ISPCState ()
compileBuiltinFun (Name
fname, func :: Function op
func@(Function Maybe EntryPoint
_ [Param]
outputs [Param]
inputs Code op
_))
  | Maybe EntryPoint -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe EntryPoint -> Bool) -> Maybe EntryPoint -> Bool
forall a b. (a -> b) -> a -> b
$ Function op -> Maybe EntryPoint
forall a. FunctionT a -> Maybe EntryPoint
functionEntry Function op
func = do
      let extra :: [Param]
extra = [[C.cparam|$tyqual:uniform struct futhark_context * $tyqual:uniform ctx|]]
          extra_c :: [Param]
extra_c = [[C.cparam|struct futhark_context * ctx|]]
          extra_exp :: [Exp]
extra_exp = [[C.cexp|$id:p|] | C.Param (Just Id
p) DeclSpec
_ Decl
_ SrcLoc
_ <- [Param]
extra]

      (inparams_c, in_args_c) <- (Param -> CompilerM Multicore ISPCState (Param, Exp))
-> [Param] -> CompilerM Multicore ISPCState ([Param], [Exp])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM ([TypeQual] -> Param -> CompilerM Multicore ISPCState (Param, Exp)
forall {op} {s}. [TypeQual] -> Param -> CompilerM op s (Param, Exp)
compileInputsExtern []) [Param]
inputs
      (outparams_c, out_args_c) <- mapAndUnzipM (compileOutputsExtern []) outputs

      (inparams_extern, _) <- mapAndUnzipM (compileInputsExtern [C.ctyquals|$tyqual:uniform|]) inputs
      (outparams_extern, _) <- mapAndUnzipM (compileOutputsExtern [C.ctyquals|$tyqual:uniform|]) outputs

      (inparams_uni, in_args_noderef) <- mapAndUnzipM compileInputsUniform inputs
      (outparams_uni, out_args_noderef) <- mapAndUnzipM compileOutputsUniform outputs

      (inparams_varying, in_args_vary, prebody_in') <- unzip3 <$> mapM compileInputsVarying inputs
      (outparams_varying, out_args_vary, prebody_out', postbody_out') <- unzip4 <$> mapM compileOutputsVarying outputs
      let (prebody_in, prebody_out, postbody_out) = over each concat (prebody_in', prebody_out', postbody_out')

      GC.libDecl
        [C.cedecl|int $id:(funName fname <> "_extern")($params:extra_c, $params:outparams_c, $params:inparams_c) {
                  return $id:(funName fname)($args:extra_exp, $args:out_args_c, $args:in_args_c);
                }|]

      let ispc_extern =
            [C.cedecl|extern "C" $tyqual:unmasked $tyqual:uniform int $id:((funName fname) <> "_extern")
                      ($params:extra, $params:outparams_extern, $params:inparams_extern);|]

          ispc_uniform =
            [C.cedecl|$tyqual:uniform int $id:(funName fname)
                    ($params:extra, $params:outparams_uni, $params:inparams_uni) {
                      return $id:(funName (fname<>"_extern"))(
                        $args:extra_exp,
                        $args:out_args_noderef,
                        $args:in_args_noderef);
                    }|]

          ispc_varying =
            [C.cedecl|$tyqual:uniform int $id:(funName fname)
                    ($params:extra, $params:outparams_varying, $params:inparams_varying) {
                        $tyqual:uniform int err = 0;
                        $items:prebody_in
                        $items:prebody_out
                        $escstm:("foreach_active (i)")
                        {
                          err |= $id:(funName $ fname<>"_extern")(
                            $args:extra_exp,
                            $args:out_args_vary,
                            $args:in_args_vary);
                        }
                        $items:postbody_out
                        return err;
                    }|]

      mapM_ ispcEarlyDecl [ispc_varying, ispc_uniform, ispc_extern]
  | Bool
otherwise = () -> CompilerM Multicore ISPCState ()
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    compileInputsExtern :: [TypeQual] -> Param -> CompilerM op s (Param, Exp)
compileInputsExtern [TypeQual]
vari (ScalarParam VName
name PrimType
bt) = do
      let ctp :: Type
ctp = PrimType -> Type
GC.primTypeToCType PrimType
bt
      (Param, Exp) -> CompilerM op s (Param, Exp)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([C.cparam|$tyquals:vari $ty:ctp $id:name|], [C.cexp|$id:name|])
    compileInputsExtern [TypeQual]
vari (MemParam VName
name Space
space) = do
      ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
GC.memToCType VName
name Space
space
      pure ([C.cparam|$tyquals:vari $ty:ty * $tyquals:vari $id:name|], [C.cexp|*$id:name|])

    compileOutputsExtern :: [TypeQual] -> Param -> CompilerM op s (Param, Exp)
compileOutputsExtern [TypeQual]
vari (ScalarParam VName
name PrimType
bt) = do
      p_name <- String -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> CompilerM op s VName) -> String -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ String
"out_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
baseString VName
name
      let ctp = PrimType -> Type
GC.primTypeToCType PrimType
bt
      pure ([C.cparam|$tyquals:vari $ty:ctp * $tyquals:vari $id:p_name|], [C.cexp|$id:p_name|])
    compileOutputsExtern [TypeQual]
vari (MemParam VName
name Space
space) = do
      ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
GC.memToCType VName
name Space
space
      p_name <- newVName $ baseString name ++ "_p"
      pure ([C.cparam|$tyquals:vari $ty:ty * $tyquals:vari $id:p_name|], [C.cexp|$id:p_name|])

    compileInputsUniform :: Param -> CompilerM op s (Param, Exp)
compileInputsUniform (ScalarParam VName
name PrimType
bt) = do
      let ctp :: Type
ctp = PrimType -> Type
GC.primTypeToCType PrimType
bt
          params :: Param
params = [C.cparam|$tyqual:uniform $ty:ctp $id:name|]
          args :: Exp
args = [C.cexp|$id:name|]
      (Param, Exp) -> CompilerM op s (Param, Exp)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param
params, Exp
args)
    compileInputsUniform (MemParam VName
name Space
space) = do
      ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
GC.memToCType VName
name Space
space
      let params = [C.cparam|$tyqual:uniform $ty:ty $id:name|]
          args = [C.cexp|&$id:name|]
      pure (params, args)

    compileOutputsUniform :: Param -> CompilerM op s (Param, Exp)
compileOutputsUniform (ScalarParam VName
name PrimType
bt) = do
      p_name <- String -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> CompilerM op s VName) -> String -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ String
"out_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
baseString VName
name
      let ctp = PrimType -> Type
GC.primTypeToCType PrimType
bt
          params = [C.cparam|$tyqual:uniform $ty:ctp *$tyqual:uniform $id:p_name|]
          args = [C.cexp|$id:p_name|]
      pure (params, args)
    compileOutputsUniform (MemParam VName
name Space
space) = do
      ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
GC.memToCType VName
name Space
space
      p_name <- newVName $ baseString name ++ "_p"
      let params = [C.cparam|$tyqual:uniform $ty:ty $id:p_name|]
          args = [C.cexp|&$id:p_name|]
      pure (params, args)

    compileInputsVarying :: Param -> CompilerM op s (Param, Exp, [BlockItem])
compileInputsVarying (ScalarParam VName
name PrimType
bt) = do
      let ctp :: Type
ctp = PrimType -> Type
GC.primTypeToCType PrimType
bt
          params :: Param
params = [C.cparam|$ty:ctp $id:name|]
          args :: Exp
args = [C.cexp|extract($id:name,i)|]
          pre_body :: [a]
pre_body = []
      (Param, Exp, [BlockItem])
-> CompilerM op s (Param, Exp, [BlockItem])
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param
params, Exp
args, [BlockItem]
forall a. [a]
pre_body)
    compileInputsVarying (MemParam VName
name Space
space) = do
      typ <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
GC.memToCType VName
name Space
space
      newvn <- newVName $ "aos_" <> baseString name
      let params = [C.cparam|$ty:typ $id:name|]
          args = [C.cexp|&$id:(newvn)[i]|]
          pre_body =
            [C.citems|$tyqual:uniform $ty:typ $id:(newvn)[programCount];
                               $id:(newvn)[programIndex] = $id:name;|]
      pure (params, args, pre_body)

    compileOutputsVarying :: Param -> CompilerM op s (Param, Exp, [BlockItem], [BlockItem])
compileOutputsVarying (ScalarParam VName
name PrimType
bt) = do
      p_name <- String -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> CompilerM op s VName) -> String -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ String
"out_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
baseString VName
name
      deref_name <- newVName $ "aos_" ++ baseString name
      vari_p_name <- newVName $ "convert_" ++ baseString name
      let ctp = PrimType -> Type
GC.primTypeToCType PrimType
bt
          pre_body =
            [C.citems|$tyqual:varying $ty:ctp $id:vari_p_name = *$id:p_name;
                                $tyqual:uniform $ty:ctp $id:deref_name[programCount];
                                $id:deref_name[programIndex] = $id:vari_p_name;|]
          post_body = [C.citems|*$id:p_name = $id:(deref_name)[programIndex];|]
          params = [C.cparam|$tyqual:varying $ty:ctp * $tyqual:uniform $id:p_name|]
          args = [C.cexp|&$id:(deref_name)[i]|]
      pure (params, args, pre_body, post_body)
    compileOutputsVarying (MemParam VName
name Space
space) = do
      typ <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
GC.memToCType VName
name Space
space
      newvn <- newVName $ "aos_" <> baseString name
      let params = [C.cparam|$ty:typ $id:name|]
          args = [C.cexp|&$id:(newvn)[i]|]
          pre_body =
            [C.citems|$tyqual:uniform $ty:typ $id:(newvn)[programCount];
                       $id:(newvn)[programIndex] = $id:name;|]
      pure (params, args, pre_body, [])

-- | Handle logging an error message in ISPC.
handleError :: ErrorMsg Exp -> String -> ISPCCompilerM ()
handleError :: ErrorMsg Exp -> String -> CompilerM Multicore ISPCState ()
handleError ErrorMsg Exp
msg String
stacktrace = do
  -- Get format sting
  (formatstr, formatargs) <- ErrorMsg Exp -> CompilerM Multicore ISPCState (String, [Exp])
forall op s. ErrorMsg Exp -> CompilerM op s (String, [Exp])
GC.errorMsgString ErrorMsg Exp
msg
  let formatstr' = String
"Error: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
formatstr String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"\n\nBacktrace:\n%s"
  -- Get args types and names for shim
  let arg_types = ErrorMsg Exp -> [PrimType]
forall a. ErrorMsg a -> [PrimType]
errorMsgArgTypes ErrorMsg Exp
msg
  arg_names <- mapM (newVName . const "arg") arg_types
  let params = (PrimType -> VName -> Param) -> [PrimType] -> [VName] -> [Param]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\PrimType
ty VName
name -> [C.cparam|$ty:(GC.primTypeToCType ty) $id:name|]) [PrimType]
arg_types [VName]
arg_names
  let params_uni = (PrimType -> VName -> Param) -> [PrimType] -> [VName] -> [Param]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\PrimType
ty VName
name -> [C.cparam|$tyqual:uniform $ty:(GC.primTypeToCType ty) $id:name|]) [PrimType]
arg_types [VName]
arg_names
  -- Make shim
  let formatargs' = ErrorMsg Exp -> [Exp] -> [VName] -> [Exp]
forall {a} {a}. ToIdent a => ErrorMsg a -> [Exp] -> [a] -> [Exp]
mapArgNames ErrorMsg Exp
msg [Exp]
formatargs [VName]
arg_names
  shim <- MC.multicoreDef "assert_shim" $ \Name
s -> do
    Definition -> CompilerM Multicore ISPCState Definition
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      [C.cedecl|void $id:s(struct futhark_context* ctx, $params:params) {
          set_error(ctx, msgprintf($string:formatstr', $args:formatargs', $string:stacktrace));
      }|]
  ispcDecl
    [C.cedecl|extern "C" $tyqual:unmasked void $id:shim($tyqual:uniform struct futhark_context* $tyqual:uniform, $params:params_uni);|]
  -- Call the shim
  args <- getErrorValExps msg
  uni <- newVName "uni"
  let args' = (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (\Exp
x -> [C.cexp|extract($exp:x, $id:uni)|]) [Exp]
args
  GC.items
    [C.citems|
      $escstm:("foreach_active(" <> prettyString uni <> ")")
      {
        $id:shim(ctx, $args:args');
        err = FUTHARK_PROGRAM_ERROR;
      }
      $escstm:("unmasked { return err; }")|]
  where
    getErrorVal :: ErrorMsgPart a -> Maybe a
getErrorVal (ErrorString Text
_) = Maybe a
forall a. Maybe a
Nothing
    getErrorVal (ErrorVal PrimType
_ a
v) = a -> Maybe a
forall a. a -> Maybe a
Just a
v

    getErrorValExps :: ErrorMsg Exp -> CompilerM Multicore ISPCState [Exp]
getErrorValExps (ErrorMsg [ErrorMsgPart Exp]
m) = (Exp -> CompilerM Multicore ISPCState Exp)
-> [Exp] -> CompilerM Multicore ISPCState [Exp]
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) -> [a] -> m [b]
mapM Exp -> CompilerM Multicore ISPCState Exp
compileExp ([Exp] -> CompilerM Multicore ISPCState [Exp])
-> [Exp] -> CompilerM Multicore ISPCState [Exp]
forall a b. (a -> b) -> a -> b
$ (ErrorMsgPart Exp -> Maybe Exp) -> [ErrorMsgPart Exp] -> [Exp]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ErrorMsgPart Exp -> Maybe Exp
forall {a}. ErrorMsgPart a -> Maybe a
getErrorVal [ErrorMsgPart Exp]
m

    mapArgNames' :: [ErrorMsgPart a] -> [Exp] -> [a] -> [Exp]
mapArgNames' (ErrorMsgPart a
x : [ErrorMsgPart a]
xs) (Exp
y : [Exp]
ys) (a
t : [a]
ts)
      | Maybe a -> Bool
forall a. Maybe a -> Bool
isJust (Maybe a -> Bool) -> Maybe a -> Bool
forall a b. (a -> b) -> a -> b
$ ErrorMsgPart a -> Maybe a
forall {a}. ErrorMsgPart a -> Maybe a
getErrorVal ErrorMsgPart a
x = [C.cexp|$id:t|] Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a] -> [Exp] -> [a] -> [Exp]
mapArgNames' [ErrorMsgPart a]
xs [Exp]
ys [a]
ts
      | Bool
otherwise = Exp
y Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a] -> [Exp] -> [a] -> [Exp]
mapArgNames' [ErrorMsgPart a]
xs [Exp]
ys (a
t a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ts)
    mapArgNames' [ErrorMsgPart a]
_ [Exp]
ys [] = [Exp]
ys
    mapArgNames' [ErrorMsgPart a]
_ [Exp]
_ [a]
_ = []

    mapArgNames :: ErrorMsg a -> [Exp] -> [a] -> [Exp]
mapArgNames (ErrorMsg [ErrorMsgPart a]
parts) = [ErrorMsgPart a] -> [Exp] -> [a] -> [Exp]
forall {a} {a}.
ToIdent a =>
[ErrorMsgPart a] -> [Exp] -> [a] -> [Exp]
mapArgNames' [ErrorMsgPart a]
parts

-- | Given the name and type of a parameter, return the C type used to
-- represent it. We use uniform pointers to varying values for lexical
-- memory blocks, as this generally results in less gathers/scatters.
getMemType :: VName -> PrimType -> ISPCCompilerM C.Type
getMemType :: VName -> PrimType -> ISPCCompilerM Type
getMemType VName
dest PrimType
elemtype = do
  cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM Multicore ISPCState (Maybe VName)
-> CompilerM Multicore ISPCState Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM Multicore ISPCState (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
GC.cacheMem VName
dest
  if cached
    then pure [C.cty|$tyqual:varying $ty:(primStorageType elemtype)* uniform|]
    else pure [C.cty|$ty:(primStorageType elemtype)*|]

compileExp :: Exp -> ISPCCompilerM C.Exp
compileExp :: Exp -> CompilerM Multicore ISPCState Exp
compileExp e :: Exp
e@(ValueExp (FloatValue (Float64Value Double
v))) =
  if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
v Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
v
    then Exp -> CompilerM Multicore ISPCState Exp
forall op s. Exp -> CompilerM op s Exp
GC.compileExp Exp
e
    else Exp -> CompilerM Multicore ISPCState Exp
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cexp|$esc:(prettyString v <> "d")|]
compileExp e :: Exp
e@(ValueExp (FloatValue (Float16Value Half
v))) =
  if Half -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Half
v Bool -> Bool -> Bool
|| Half -> Bool
forall a. RealFloat a => a -> Bool
isNaN Half
v
    then Exp -> CompilerM Multicore ISPCState Exp
forall op s. Exp -> CompilerM op s Exp
GC.compileExp Exp
e
    else Exp -> CompilerM Multicore ISPCState Exp
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cexp|$esc:(prettyString v <> "f16")|]
compileExp (ValueExp PrimValue
val) =
  Exp -> CompilerM Multicore ISPCState Exp
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> CompilerM Multicore ISPCState Exp)
-> Exp -> CompilerM Multicore ISPCState Exp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
C.toExp PrimValue
val SrcLoc
forall a. Monoid a => a
mempty
compileExp (LeafExp VName
v PrimType
_) =
  Exp -> CompilerM Multicore ISPCState Exp
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cexp|$id:v|]
compileExp (UnOpExp Complement {} Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|~$exp:x'|]
compileExp (UnOpExp (Neg PrimType
Bool) Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|!$exp:x'|]
compileExp (UnOpExp Neg {} Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|-$exp:x'|]
compileExp (UnOpExp (FAbs FloatType
Float32) Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|(float)fabs($exp:x')|]
compileExp (UnOpExp (FAbs FloatType
Float64) Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|fabs($exp:x')|]
compileExp (UnOpExp SSignum {} Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|($exp:x' > 0 ? 1 : 0) - ($exp:x' < 0 ? 1 : 0)|]
compileExp (UnOpExp USignum {} Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|($exp:x' > 0 ? 1 : 0) - ($exp:x' < 0 ? 1 : 0) != 0|]
compileExp (UnOpExp UnOp
op Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|$id:(prettyString op)($exp:x')|]
compileExp (CmpOpExp CmpOp
cmp Exp
x Exp
y) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  y' <- compileExp y
  pure $ case cmp of
    CmpEq {} -> [C.cexp|$exp:x' == $exp:y'|]
    FCmpLt {} -> [C.cexp|$exp:x' < $exp:y'|]
    FCmpLe {} -> [C.cexp|$exp:x' <= $exp:y'|]
    CmpLlt {} -> [C.cexp|$exp:x' < $exp:y'|]
    CmpLle {} -> [C.cexp|$exp:x' <= $exp:y'|]
    CmpOp
_ -> [C.cexp|$id:(prettyString cmp)($exp:x', $exp:y')|]
compileExp (ConvOpExp ConvOp
conv Exp
x) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  pure [C.cexp|$id:(prettyString conv)($exp:x')|]
compileExp (BinOpExp BinOp
bop Exp
x Exp
y) = do
  x' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
x
  y' <- compileExp y
  pure $ case bop of
    Add IntType
_ Overflow
OverflowUndef -> [C.cexp|$exp:x' + $exp:y'|]
    Sub IntType
_ Overflow
OverflowUndef -> [C.cexp|$exp:x' - $exp:y'|]
    Mul IntType
_ Overflow
OverflowUndef -> [C.cexp|$exp:x' * $exp:y'|]
    FAdd {} -> [C.cexp|$exp:x' + $exp:y'|]
    FSub {} -> [C.cexp|$exp:x' - $exp:y'|]
    FMul {} -> [C.cexp|$exp:x' * $exp:y'|]
    FDiv {} -> [C.cexp|$exp:x' / $exp:y'|]
    Xor {} -> [C.cexp|$exp:x' ^ $exp:y'|]
    And {} -> [C.cexp|$exp:x' & $exp:y'|]
    Or {} -> [C.cexp|$exp:x' | $exp:y'|]
    LogAnd {} -> [C.cexp|$exp:x' && $exp:y'|]
    LogOr {} -> [C.cexp|$exp:x' || $exp:y'|]
    BinOp
_ -> [C.cexp|$id:(prettyString bop)($exp:x', $exp:y')|]
compileExp (FunExp Text
h [Exp]
args PrimType
_) = do
  args' <- (Exp -> CompilerM Multicore ISPCState Exp)
-> [Exp] -> CompilerM Multicore ISPCState [Exp]
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) -> [a] -> m [b]
mapM Exp -> CompilerM Multicore ISPCState Exp
compileExp [Exp]
args
  pure [C.cexp|$id:(funName (nameFromText h))($args:args')|]

-- | Compile a block of code with ISPC specific semantics, falling back
-- to generic C when this semantics is not needed.
-- All recursive constructors are duplicated here, since not doing so
-- would cause use to enter regular generic C codegen with no escape.
compileCode :: MCCode -> ISPCCompilerM ()
compileCode :: MCCode -> CompilerM Multicore ISPCState ()
compileCode (Meta (MetaComment Text
s)) = do
  Text -> CompilerM Multicore ISPCState ()
forall op s. Text -> CompilerM op s ()
GC.comment Text
s
compileCode (DeclareScalar VName
name Volatility
_ PrimType
t) = do
  let ct :: Type
ct = PrimType -> Type
GC.primTypeToCType PrimType
t
  quals <- VName -> ISPCCompilerM [TypeQual]
getVariabilityQuals VName
name
  GC.decl [C.cdecl|$tyquals:quals $ty:ct $id:name;|]
compileCode (DeclareArray VName
name PrimType
t ArrayContents
vs) = do
  name_realtype <- String -> CompilerM Multicore ISPCState VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> CompilerM Multicore ISPCState VName)
-> String -> CompilerM Multicore ISPCState VName
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString VName
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_realtype"
  let ct = PrimType -> Type
GC.primTypeToCType PrimType
t
  case vs of
    ArrayValues [PrimValue]
vs' -> do
      let vs'' :: [Initializer]
vs'' = [[C.cinit|$exp:v|] | PrimValue
v <- [PrimValue]
vs']
      Definition -> CompilerM Multicore ISPCState ()
forall op s. Definition -> CompilerM op s ()
GC.earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs')] = {$inits:vs''};|]
    ArrayZeros Int
n ->
      Definition -> CompilerM Multicore ISPCState ()
forall op s. Definition -> CompilerM op s ()
GC.earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|]
  -- Make an exported C shim to access a faked memory block.
  shim <- MC.multicoreDef "get_static_array_shim" $ \Name
f ->
    Definition -> CompilerM Multicore ISPCState Definition
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      [C.cedecl|struct memblock $id:f(struct futhark_context* ctx) {
                  return (struct memblock){NULL,(unsigned char*)$id:name_realtype,0};
                }|]
  ispcDecl
    [C.cedecl|extern "C" $tyqual:unmasked $tyqual:uniform struct memblock $tyqual:uniform
                        $id:shim($tyqual:uniform struct futhark_context* $tyqual:uniform ctx);|]
  -- Call it
  GC.item [C.citem|$tyqual:uniform struct memblock $id:name = $id:shim(ctx);|]
compileCode (MCCode
c1 :>>: MCCode
c2) = [MCCode] -> CompilerM Multicore ISPCState ()
go (MCCode -> [MCCode]
forall op. Code op -> [Code op]
GC.linearCode (MCCode
c1 MCCode -> MCCode -> MCCode
forall a. Code a -> Code a -> Code a
:>>: MCCode
c2))
  where
    go :: [MCCode] -> CompilerM Multicore ISPCState ()
go (DeclareScalar VName
name Volatility
_ PrimType
t : SetScalar VName
dest Exp
e : [MCCode]
code)
      | VName
name VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
dest = do
          let ct :: Type
ct = PrimType -> Type
GC.primTypeToCType PrimType
t
          e' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
e
          quals <- getVariabilityQuals name
          GC.item [C.citem|$tyquals:quals $ty:ct $id:name = $exp:e';|]
          go code
    go (MCCode
x : [MCCode]
xs) = MCCode -> CompilerM Multicore ISPCState ()
compileCode MCCode
x CompilerM Multicore ISPCState ()
-> CompilerM Multicore ISPCState ()
-> CompilerM Multicore ISPCState ()
forall a b.
CompilerM Multicore ISPCState a
-> CompilerM Multicore ISPCState b
-> CompilerM Multicore ISPCState b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [MCCode] -> CompilerM Multicore ISPCState ()
go [MCCode]
xs
    go [] = () -> CompilerM Multicore ISPCState ()
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
compileCode (Allocate VName
name (Count (TPrimExp Exp
e)) Space
space) = do
  size <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
e
  cached <- GC.cacheMem name
  case cached of
    Just VName
cur_size ->
      Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm
        [C.cstm|if ($exp:cur_size < $exp:size) {
                  err = lexical_realloc(ctx, &$exp:name, &$exp:cur_size, $exp:size);
                  if (err != FUTHARK_SUCCESS) {
                    $escstm:("unmasked { return err; }")
                  }
                }|]
    Maybe VName
_ ->
      VName -> Exp -> Space -> Stm -> CompilerM Multicore ISPCState ()
forall a b.
(ToExp a, ToExp b) =>
a -> b -> Space -> Stm -> CompilerM Multicore ISPCState ()
allocMem VName
name Exp
size Space
space [C.cstm|$escstm:("unmasked { return 1; }")|]
compileCode (SetMem VName
dest VName
src Space
space) =
  VName -> VName -> Space -> CompilerM Multicore ISPCState ()
forall a b.
(ToExp a, ToExp b) =>
a -> b -> Space -> CompilerM Multicore ISPCState ()
setMem VName
dest VName
src Space
space
compileCode (Write VName
dest (Count TExp Int64
idx) PrimType
elemtype Space
DefaultSpace Volatility
_ Exp
elemexp)
  | Exp -> Bool
forall {v}. PrimExp v -> Bool
isConstExp (TExp Int64 -> Exp
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TExp Int64
idx) = do
      dest' <- VName -> CompilerM Multicore ISPCState Exp
forall op s. VName -> CompilerM op s Exp
GC.rawMem VName
dest
      idxexp <- compileExp $ constFoldPrimExp $ untyped idx
      deref <-
        GC.derefPointer
          dest'
          [C.cexp|($tyquals:([varying]) typename int64_t)$exp:idxexp|]
          <$> getMemType dest elemtype
      elemexp' <- toStorage elemtype <$> compileExp elemexp
      GC.stm [C.cstm|$exp:deref = $exp:elemexp';|]
  | Bool
otherwise = do
      dest' <- VName -> CompilerM Multicore ISPCState Exp
forall op s. VName -> CompilerM op s Exp
GC.rawMem VName
dest
      idxexp <- compileExp $ untyped idx
      deref <-
        GC.derefPointer
          dest'
          [C.cexp|($tyquals:([varying]) typename int64_t)$exp:idxexp|]
          <$> getMemType dest elemtype
      elemexp' <- toStorage elemtype <$> compileExp elemexp
      GC.stm [C.cstm|$exp:deref = $exp:elemexp';|]
  where
    isConstExp :: PrimExp v -> Bool
isConstExp = PrimExp v -> Bool
forall {v}. PrimExp v -> Bool
isSimple (PrimExp v -> Bool)
-> (PrimExp v -> PrimExp v) -> PrimExp v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimExp v -> PrimExp v
forall v. PrimExp v -> PrimExp v
constFoldPrimExp
    isSimple :: PrimExp v -> Bool
isSimple (ValueExp PrimValue
_) = Bool
True
    isSimple PrimExp v
_ = Bool
False
compileCode (Read VName
x VName
src (Count TExp Int64
iexp) PrimType
restype Space
DefaultSpace Volatility
_) = do
  src' <- VName -> CompilerM Multicore ISPCState Exp
forall op s. VName -> CompilerM op s Exp
GC.rawMem VName
src
  e <-
    fmap (fromStorage restype) $
      GC.derefPointer src'
        <$> compileExp (untyped iexp)
        <*> getMemType src restype
  GC.stm [C.cstm|$id:x = $exp:e;|]
compileCode (Copy PrimType
t [Count Elements (TExp Int64)]
shape (VName
dst, Space
DefaultSpace) (Count Elements (TExp Int64), [Count Elements (TExp Int64)])
dst_lmad (VName
src, Space
DefaultSpace) (Count Elements (TExp Int64), [Count Elements (TExp Int64)])
src_lmad) = do
  dst' <- VName -> CompilerM Multicore ISPCState Exp
forall op s. VName -> CompilerM op s Exp
GC.rawMem VName
dst
  src' <- GC.rawMem src
  let doWrite a
dst_i Exp
ve = do
        deref <-
          Exp -> Exp -> Type -> Exp
GC.derefPointer
            Exp
dst'
            [C.cexp|($tyquals:([varying]) typename int64_t)$exp:dst_i|]
            (Type -> Exp)
-> ISPCCompilerM Type -> CompilerM Multicore ISPCState Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> PrimType -> ISPCCompilerM Type
getMemType VName
dst PrimType
t
        GC.stm [C.cstm|$exp:deref = $exp:(toStorage t ve);|]
      doRead Exp
src_i =
        PrimType -> Exp -> Exp
fromStorage PrimType
t (Exp -> Exp) -> (Type -> Exp) -> Type -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Exp -> Type -> Exp
GC.derefPointer Exp
src' Exp
src_i (Type -> Exp)
-> ISPCCompilerM Type -> CompilerM Multicore ISPCState Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> PrimType -> ISPCCompilerM Type
getMemType VName
src PrimType
t
  GC.compileCopyWith shape doWrite dst_lmad doRead src_lmad
compileCode (Free VName
name Space
space) = do
  cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM Multicore ISPCState (Maybe VName)
-> CompilerM Multicore ISPCState Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM Multicore ISPCState (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
GC.cacheMem VName
name
  unless cached $ unRefMem name space
compileCode (For VName
i Exp
bound MCCode
body)
  -- The special-case here is to avoid certain pathological/contrived
  -- programs that construct statically known zero-element arrays.
  -- Due to the way we do constant-fold index functions, this produces
  -- code that looks like it has uniform/varying mismatches (i.e. race
  -- conditions) to ISPC, even though that code is never actually run.
  | Exp -> Bool
forall {v}. PrimExp v -> Bool
isZero Exp
bound = () -> CompilerM Multicore ISPCState ()
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise = do
      let i' :: SrcLoc -> Id
i' = VName -> SrcLoc -> Id
forall a. ToIdent a => a -> SrcLoc -> Id
C.toIdent VName
i
          t :: Type
t = PrimType -> Type
GC.primTypeToCType (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ Exp -> PrimType
forall v. PrimExp v -> PrimType
primExpType Exp
bound
      bound' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
bound
      body' <- GC.collect $ compileCode body
      quals <- getVariabilityQuals i
      GC.stm
        [C.cstm|for ($tyquals:quals $ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) {
            $items:body'
          }|]
  where
    isZero :: PrimExp v -> Bool
isZero (ValueExp PrimValue
v) = PrimValue -> Bool
zeroIsh PrimValue
v
    isZero PrimExp v
_ = Bool
False
compileCode (While TExp Bool
cond MCCode
body) = do
  cond' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp (Exp -> CompilerM Multicore ISPCState Exp)
-> Exp -> CompilerM Multicore ISPCState Exp
forall a b. (a -> b) -> a -> b
$ TExp Bool -> Exp
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TExp Bool
cond
  body' <- GC.collect $ compileCode body
  GC.stm
    [C.cstm|while ($exp:cond') {
            $items:body'
          }|]
compileCode (If TExp Bool
cond MCCode
tbranch MCCode
fbranch) = do
  cond' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp (Exp -> CompilerM Multicore ISPCState Exp)
-> Exp -> CompilerM Multicore ISPCState Exp
forall a b. (a -> b) -> a -> b
$ TExp Bool -> Exp
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TExp Bool
cond
  tbranch' <- GC.collect $ compileCode tbranch
  fbranch' <- GC.collect $ compileCode fbranch
  GC.stm $ case (tbranch', fbranch') of
    ([BlockItem]
_, []) ->
      [C.cstm|if ($exp:cond') { $items:tbranch' }|]
    ([], [BlockItem]
_) ->
      [C.cstm|if (!($exp:cond')) { $items:fbranch' }|]
    ([BlockItem], [BlockItem])
_ ->
      [C.cstm|if ($exp:cond') { $items:tbranch' } else { $items:fbranch' }|]
compileCode (Call [VName]
dests Name
fname [Arg]
args) = do
  (dests', unpack_dest) <- (VName -> CompilerM Multicore ISPCState (VName, [Stm]))
-> [VName] -> CompilerM Multicore ISPCState ([VName], [[Stm]])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM VName -> CompilerM Multicore ISPCState (VName, [Stm])
forall op s. VName -> CompilerM op s (VName, [Stm])
GC.compileDest [VName]
dests
  defCallIspc dests' fname =<< mapM GC.compileArg args
  GC.stms $ mconcat unpack_dest
  where
    defCallIspc :: [a] -> Name -> [Exp] -> CompilerM op s ()
defCallIspc [a]
dests' Name
fname' [Exp]
args' = do
      let out_args :: [Exp]
out_args = [[C.cexp|&$id:d|] | a
d <- [a]
dests']
          args'' :: [Exp]
args''
            | Name -> Bool
isBuiltInFunction Name
fname' = [Exp]
args'
            | Bool
otherwise = [C.cexp|ctx|] Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp]
out_args [Exp] -> [Exp] -> [Exp]
forall a. [a] -> [a] -> [a]
++ [Exp]
args'
      case [a]
dests' of
        [a
d]
          | Name -> Bool
isBuiltInFunction Name
fname' ->
              Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:d = $id:(funName fname')($args:args'');|]
        [a]
_ ->
          BlockItem -> CompilerM op s ()
forall op s. BlockItem -> CompilerM op s ()
GC.item
            [C.citem|
            if ($id:(funName fname')($args:args'') != 0) {
              $escstm:("unmasked { return 1; }")
            }|]
compileCode (Assert Exp
e ErrorMsg Exp
msg (Loc
loc, [Loc]
locs)) = do
  e' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
e
  err <- GC.collect $ handleError msg stacktrace
  GC.stm [C.cstm|if (!$exp:e') { $items:err }|]
  where
    stacktrace :: String
stacktrace = Text -> String
T.unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ Int -> [Text] -> Text
prettyStacktrace Int
0 ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (Loc -> Text) -> [Loc] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Loc -> Text
forall a. Located a => a -> Text
locText ([Loc] -> [Text]) -> [Loc] -> [Text]
forall a b. (a -> b) -> a -> b
$ Loc
loc Loc -> [Loc] -> [Loc]
forall a. a -> [a] -> [a]
: [Loc]
locs
compileCode MCCode
code =
  MCCode -> CompilerM Multicore ISPCState ()
forall op s. Code op -> CompilerM op s ()
GC.compileCode MCCode
code

-- | Prepare a struct with memory allocted in the scope and populate
-- its fields with values
prepareMemStruct :: [(VName, VName)] -> [VName] -> ISPCCompilerM Name
prepareMemStruct :: [(VName, VName)] -> [VName] -> CompilerM Multicore ISPCState Name
prepareMemStruct [(VName, VName)]
lexmems [VName]
fatmems = do
  let lex_defs :: [FieldGroup]
lex_defs = ((VName, VName) -> [FieldGroup])
-> [(VName, VName)] -> [FieldGroup]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (VName, VName) -> [FieldGroup]
forall {a} {a}. (ToIdent a, ToIdent a) => (a, a) -> [FieldGroup]
lexMemDef [(VName, VName)]
lexmems
  let fat_defs :: [FieldGroup]
fat_defs = (VName -> FieldGroup) -> [VName] -> [FieldGroup]
forall a b. (a -> b) -> [a] -> [b]
map VName -> FieldGroup
forall {a}. ToIdent a => a -> FieldGroup
fatMemDef [VName]
fatmems
  name <- DefSpecifier ISPCState
ispcDef String
"mem_struct" ((Name -> CompilerM Multicore ISPCState Definition)
 -> CompilerM Multicore ISPCState Name)
-> (Name -> CompilerM Multicore ISPCState Definition)
-> CompilerM Multicore ISPCState Name
forall a b. (a -> b) -> a -> b
$ \Name
s -> do
    Definition -> CompilerM Multicore ISPCState Definition
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      [C.cedecl|struct $id:s {
        $sdecls:lex_defs
        $sdecls:fat_defs
      };|]
  let name' = Name
name Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_"
  GC.decl [C.cdecl|$tyqual:uniform struct $id:name $id:name';|]
  forM_ (concatMap (\(VName
a, VName
b) -> [VName
a, VName
b]) lexmems) $ \VName
m ->
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:name'.$id:m = $id:m;|]
  forM_ fatmems $ \VName
m ->
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:name'.$id:m = &$id:m;|]
  pure name
  where
    lexMemDef :: (a, a) -> [FieldGroup]
lexMemDef (a
name, a
size) =
      [ [C.csdecl|$tyqual:varying unsigned char * $tyqual:uniform $id:name;|],
        [C.csdecl|$tyqual:varying size_t $id:size;|]
      ]
    fatMemDef :: a -> FieldGroup
fatMemDef a
name =
      [C.csdecl|$tyqual:varying struct memblock * $tyqual:uniform $id:name;|]

-- | Get memory from the memory struct into local variables
compileGetMemStructVals :: Name -> [(VName, VName)] -> [VName] -> ISPCCompilerM ()
compileGetMemStructVals :: Name
-> [(VName, VName)] -> [VName] -> CompilerM Multicore ISPCState ()
compileGetMemStructVals Name
struct [(VName, VName)]
lexmems [VName]
fatmems = do
  [VName]
-> (VName -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [VName]
fatmems ((VName -> CompilerM Multicore ISPCState ())
 -> CompilerM Multicore ISPCState ())
-> (VName -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall a b. (a -> b) -> a -> b
$ \VName
m ->
    InitGroup -> CompilerM Multicore ISPCState ()
forall op s. InitGroup -> CompilerM op s ()
GC.decl [C.cdecl|struct memblock $id:m = *$id:struct->$id:m;|]
  [(VName, VName)]
-> ((VName, VName) -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(VName, VName)]
lexmems (((VName, VName) -> CompilerM Multicore ISPCState ())
 -> CompilerM Multicore ISPCState ())
-> ((VName, VName) -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall a b. (a -> b) -> a -> b
$ \(VName
m, VName
s) -> do
    InitGroup -> CompilerM Multicore ISPCState ()
forall op s. InitGroup -> CompilerM op s ()
GC.decl [C.cdecl|$tyqual:varying unsigned char * $tyqual:uniform $id:m = $id:struct->$id:m;|]
    InitGroup -> CompilerM Multicore ISPCState ()
forall op s. InitGroup -> CompilerM op s ()
GC.decl [C.cdecl|size_t $id:s = $id:struct->$id:s;|]

-- | Write back potentially changed memory addresses and sizes to the memory struct
compileWritebackMemStructVals :: Name -> [(VName, VName)] -> [VName] -> ISPCCompilerM ()
compileWritebackMemStructVals :: Name
-> [(VName, VName)] -> [VName] -> CompilerM Multicore ISPCState ()
compileWritebackMemStructVals Name
struct [(VName, VName)]
lexmems [VName]
fatmems = do
  [VName]
-> (VName -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [VName]
fatmems ((VName -> CompilerM Multicore ISPCState ())
 -> CompilerM Multicore ISPCState ())
-> (VName -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall a b. (a -> b) -> a -> b
$ \VName
m ->
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|*$id:struct->$id:m = $id:m;|]
  [(VName, VName)]
-> ((VName, VName) -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(VName, VName)]
lexmems (((VName, VName) -> CompilerM Multicore ISPCState ())
 -> CompilerM Multicore ISPCState ())
-> ((VName, VName) -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall a b. (a -> b) -> a -> b
$ \(VName
m, VName
s) -> do
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:struct->$id:m = $id:m;|]
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:struct->$id:s = $id:s;|]

-- | Read back potentially changed memory addresses and sizes to the memory struct into local variables
compileReadbackMemStructVals :: Name -> [(VName, VName)] -> [VName] -> ISPCCompilerM ()
compileReadbackMemStructVals :: Name
-> [(VName, VName)] -> [VName] -> CompilerM Multicore ISPCState ()
compileReadbackMemStructVals Name
struct [(VName, VName)]
lexmems [VName]
fatmems = do
  [VName]
-> (VName -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [VName]
fatmems ((VName -> CompilerM Multicore ISPCState ())
 -> CompilerM Multicore ISPCState ())
-> (VName -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall a b. (a -> b) -> a -> b
$ \VName
m ->
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:m = *$id:struct.$id:m;|]
  [(VName, VName)]
-> ((VName, VName) -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(VName, VName)]
lexmems (((VName, VName) -> CompilerM Multicore ISPCState ())
 -> CompilerM Multicore ISPCState ())
-> ((VName, VName) -> CompilerM Multicore ISPCState ())
-> CompilerM Multicore ISPCState ()
forall a b. (a -> b) -> a -> b
$ \(VName
m, VName
s) -> do
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:m = $id:struct.$id:m;|]
    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:s = $id:struct.$id:s;|]

compileGetStructVals ::
  Name ->
  [VName] ->
  [(C.Type, MC.ValueType)] ->
  ISPCCompilerM [C.BlockItem]
compileGetStructVals :: Name -> [VName] -> [(Type, ValueType)] -> ISPCCompilerM [BlockItem]
compileGetStructVals Name
struct [VName]
a [(Type, ValueType)]
b = [[BlockItem]] -> [BlockItem]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[BlockItem]] -> [BlockItem])
-> CompilerM Multicore ISPCState [[BlockItem]]
-> ISPCCompilerM [BlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> (Type, ValueType) -> ISPCCompilerM [BlockItem])
-> [VName]
-> [(Type, ValueType)]
-> CompilerM Multicore ISPCState [[BlockItem]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM VName -> (Type, ValueType) -> ISPCCompilerM [BlockItem]
field [VName]
a [(Type, ValueType)]
b
  where
    struct' :: Name
struct' = Name
struct Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_"
    field :: VName -> (Type, ValueType) -> ISPCCompilerM [BlockItem]
field VName
name (Type
ty, MC.Prim PrimType
pt) = do
      let inner :: Exp
inner = [C.cexp|$id:struct'->$id:(MC.closureFreeStructField name)|]
      [BlockItem] -> ISPCCompilerM [BlockItem]
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.citems|$tyqual:uniform $ty:ty $id:name = $exp:(fromStorage pt inner);|]
    field VName
name (Type
_, ValueType
_) = do
      strlit <- String -> CompilerM Multicore ISPCState Name
makeStringLiteral (String -> CompilerM Multicore ISPCState Name)
-> String -> CompilerM Multicore ISPCState Name
forall a b. (a -> b) -> a -> b
$ VName -> String
forall a. Pretty a => a -> String
prettyString VName
name
      pure
        [C.citems|$tyqual:uniform struct memblock $id:name;
                     $id:name.desc = $id:strlit();
                     $id:name.mem = $id:struct'->$id:(MC.closureFreeStructField name);
                     $id:name.size = 0;
                     $id:name.references = NULL;|]

-- | Can the given code produce an error? If so, we can't use foreach
-- loops, since they don't allow for early-outs in error handling.
mayProduceError :: MCCode -> Bool
mayProduceError :: MCCode -> Bool
mayProduceError (MCCode
x :>>: MCCode
y) = MCCode -> Bool
mayProduceError MCCode
x Bool -> Bool -> Bool
|| MCCode -> Bool
mayProduceError MCCode
y
mayProduceError (If TExp Bool
_ MCCode
x MCCode
y) = MCCode -> Bool
mayProduceError MCCode
x Bool -> Bool -> Bool
|| MCCode -> Bool
mayProduceError MCCode
y
mayProduceError (For VName
_ Exp
_ MCCode
x) = MCCode -> Bool
mayProduceError MCCode
x
mayProduceError (While TExp Bool
_ MCCode
x) = MCCode -> Bool
mayProduceError MCCode
x
mayProduceError (Op (ForEachActive VName
_ MCCode
body)) = MCCode -> Bool
mayProduceError MCCode
body
mayProduceError (Op (ForEach VName
_ Exp
_ Exp
_ MCCode
body)) = MCCode -> Bool
mayProduceError MCCode
body
mayProduceError (Op SegOp {}) = Bool
True
mayProduceError Allocate {} = Bool
True
mayProduceError Assert {} = Bool
True
mayProduceError SetMem {} = Bool
True
mayProduceError Free {} = Bool
True
mayProduceError Call {} = Bool
True
mayProduceError MCCode
_ = Bool
False

-- Generate a segop function for top_level and potentially nested SegOp code
compileOp :: GC.OpCompiler Multicore ISPCState
compileOp :: OpCompiler Multicore ISPCState
compileOp (SegOp String
name [Param]
params ParallelTask
seq_task Maybe ParallelTask
par_task [Param]
retvals (SchedulerInfo Exp
e Scheduling
sched)) = do
  let (ParallelTask MCCode
seq_code) = ParallelTask
seq_task
  free_ctypes <- (Param -> CompilerM Multicore ISPCState (Type, ValueType))
-> [Param] -> CompilerM Multicore ISPCState [(Type, ValueType)]
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) -> [a] -> m [b]
mapM Param -> CompilerM Multicore ISPCState (Type, ValueType)
forall op s. Param -> CompilerM op s (Type, ValueType)
MC.paramToCType [Param]
params
  retval_ctypes <- mapM MC.paramToCType retvals
  let free_args = (Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
params
      retval_args = (Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
retvals
      free = [VName] -> [(Type, ValueType)] -> [(VName, (Type, ValueType))]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
free_args [(Type, ValueType)]
free_ctypes
      retval = [VName] -> [(Type, ValueType)] -> [(VName, (Type, ValueType))]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
retval_args [(Type, ValueType)]
retval_ctypes

  e' <- compileExp e

  let lexical = KernelHandling -> Function Multicore -> Map VName Space
lexicalMemoryUsageMC KernelHandling
OpaqueKernels (Function Multicore -> Map VName Space)
-> Function Multicore -> Map VName Space
forall a b. (a -> b) -> a -> b
$ Maybe EntryPoint
-> [Param] -> [Param] -> MCCode -> Function Multicore
forall a.
Maybe EntryPoint -> [Param] -> [Param] -> Code a -> FunctionT a
Function Maybe EntryPoint
forall a. Maybe a
Nothing [] [Param]
params MCCode
seq_code

  fstruct <-
    MC.prepareTaskStruct sharedDef "task" free_args free_ctypes retval_args retval_ctypes

  fpar_task <- MC.generateParLoopFn lexical (name ++ "_task") seq_code fstruct free retval
  MC.addTimingFields fpar_task

  let ftask_name = Name
fstruct Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_task"

  to_c <- GC.collect $ do
    GC.decl [C.cdecl|struct scheduler_segop $id:ftask_name;|]
    GC.stm [C.cstm|$id:ftask_name.args = args;|]
    GC.stm [C.cstm|$id:ftask_name.top_level_fn = $id:fpar_task;|]
    GC.stm [C.cstm|$id:ftask_name.name = $string:(nameToString fpar_task);|]
    GC.stm [C.cstm|$id:ftask_name.iterations = iterations;|]
    -- Create the timing fields for the task
    GC.stm [C.cstm|$id:ftask_name.task_time = &ctx->program->$id:(MC.functionTiming fpar_task);|]
    GC.stm [C.cstm|$id:ftask_name.task_iter = &ctx->program->$id:(MC.functionIterations fpar_task);|]

    case sched of
      Scheduling
Dynamic -> Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:ftask_name.sched = DYNAMIC;|]
      Scheduling
Static -> Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:ftask_name.sched = STATIC;|]

    -- Generate the nested segop function if available
    case par_task of
      Just (ParallelTask MCCode
nested_code) -> do
        let lexical_nested :: Map VName Space
lexical_nested = KernelHandling -> Function Multicore -> Map VName Space
lexicalMemoryUsageMC KernelHandling
OpaqueKernels (Function Multicore -> Map VName Space)
-> Function Multicore -> Map VName Space
forall a b. (a -> b) -> a -> b
$ Maybe EntryPoint
-> [Param] -> [Param] -> MCCode -> Function Multicore
forall a.
Maybe EntryPoint -> [Param] -> [Param] -> Code a -> FunctionT a
Function Maybe EntryPoint
forall a. Maybe a
Nothing [] [Param]
params MCCode
nested_code
        fnpar_task <- Map VName Space
-> String
-> MCCode
-> Name
-> [(VName, (Type, ValueType))]
-> [(VName, (Type, ValueType))]
-> CompilerM Multicore ISPCState Name
forall a s.
ToIdent a =>
Map VName Space
-> String
-> MCCode
-> a
-> [(VName, (Type, ValueType))]
-> [(VName, (Type, ValueType))]
-> CompilerM Multicore s Name
MC.generateParLoopFn Map VName Space
lexical_nested (String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_nested_task") MCCode
nested_code Name
fstruct [(VName, (Type, ValueType))]
free [(VName, (Type, ValueType))]
retval
        GC.stm [C.cstm|$id:ftask_name.nested_fn = $id:fnpar_task;|]
      Maybe ParallelTask
Nothing ->
        Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|]

    GC.stm [C.cstm|return scheduler_prepare_task(&ctx->scheduler, &$id:ftask_name);|]

  schedn <- MC.multicoreDef "schedule_shim" $ \Name
s ->
    Definition -> CompilerM Multicore ISPCState Definition
forall a. a -> CompilerM Multicore ISPCState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      [C.cedecl|int $id:s(struct futhark_context* ctx, void* args, typename int64_t iterations) {
        $items:to_c
    }|]

  ispcDecl
    [C.cedecl|extern "C" $tyqual:unmasked $tyqual:uniform int $id:schedn
                        (struct futhark_context $tyqual:uniform * $tyqual:uniform ctx,
                        struct $id:fstruct $tyqual:uniform * $tyqual:uniform args,
                        $tyqual:uniform int iterations);|]

  aos_name <- newVName "aos"
  GC.items
    [C.citems|
    $escstm:("#if defined(ISPC)")
    $tyqual:uniform struct $id:fstruct $id:aos_name[programCount];
    $id:aos_name[programIndex] = $id:(fstruct <> "_");
    $escstm:("foreach_active (i)")
    {
      if (err == 0) {
        err = $id:schedn(ctx, &$id:aos_name[i], extract($exp:e', i));
      }
    }
    if (err != 0) {
      $escstm:("unmasked { return err; }")
    }
    $escstm:("#else")
    err = $id:schedn(ctx, &$id:(fstruct <> "_"), $exp:e');
    if (err != 0) {
      goto cleanup;
    }
    $escstm:("#endif")|]
compileOp (ISPCKernel MCCode
body [Param]
free) = do
  free_ctypes <- (Param -> CompilerM Multicore ISPCState (Type, ValueType))
-> [Param] -> CompilerM Multicore ISPCState [(Type, ValueType)]
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) -> [a] -> m [b]
mapM Param -> CompilerM Multicore ISPCState (Type, ValueType)
forall op s. Param -> CompilerM op s (Type, ValueType)
MC.paramToCType [Param]
free
  let free_args = (Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
free

  let lexical = KernelHandling -> Function Multicore -> Map VName Space
lexicalMemoryUsageMC KernelHandling
OpaqueKernels (Function Multicore -> Map VName Space)
-> Function Multicore -> Map VName Space
forall a b. (a -> b) -> a -> b
$ Maybe EntryPoint
-> [Param] -> [Param] -> MCCode -> Function Multicore
forall a.
Maybe EntryPoint -> [Param] -> [Param] -> Code a -> FunctionT a
Function Maybe EntryPoint
forall a. Maybe a
Nothing [] [Param]
free MCCode
body
  -- Generate ISPC kernel
  fstruct <- MC.prepareTaskStruct sharedDef "param_struct" free_args free_ctypes [] []
  let fstruct' = Name
fstruct Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_"

  ispcShim <- ispcDef "loop_ispc" $ \Name
s -> do
    mainBody <- ISPCCompilerM [BlockItem] -> ISPCCompilerM [BlockItem]
forall op s a. CompilerM op s a -> CompilerM op s a
GC.inNewFunction (ISPCCompilerM [BlockItem] -> ISPCCompilerM [BlockItem])
-> ISPCCompilerM [BlockItem] -> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$
      MCCode -> ISPCCompilerM [BlockItem] -> ISPCCompilerM [BlockItem]
forall a. MCCode -> ISPCCompilerM a -> ISPCCompilerM a
analyzeVariability MCCode
body (ISPCCompilerM [BlockItem] -> ISPCCompilerM [BlockItem])
-> ISPCCompilerM [BlockItem] -> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$
        Map VName Space
-> ([BlockItem]
    -> [Stm] -> [(VName, VName)] -> ISPCCompilerM [BlockItem])
-> ISPCCompilerM [BlockItem]
forall op s a.
Map VName Space
-> ([BlockItem] -> [Stm] -> [(VName, VName)] -> CompilerM op s a)
-> CompilerM op s a
cachingMemory Map VName Space
lexical (([BlockItem]
  -> [Stm] -> [(VName, VName)] -> ISPCCompilerM [BlockItem])
 -> ISPCCompilerM [BlockItem])
-> ([BlockItem]
    -> [Stm] -> [(VName, VName)] -> ISPCCompilerM [BlockItem])
-> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$ \[BlockItem]
decl_cached [Stm]
free_cached [(VName, VName)]
lexmems ->
          CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
GC.collect (CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem])
-> CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$ do
            InitGroup -> CompilerM Multicore ISPCState ()
forall op s. InitGroup -> CompilerM op s ()
GC.decl [C.cdecl|$tyqual:uniform struct futhark_context * $tyqual:uniform ctx = $id:fstruct'->ctx;|]
            [BlockItem] -> CompilerM Multicore ISPCState ()
forall op s. [BlockItem] -> CompilerM op s ()
GC.items ([BlockItem] -> CompilerM Multicore ISPCState ())
-> ISPCCompilerM [BlockItem] -> CompilerM Multicore ISPCState ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> [VName] -> [(Type, ValueType)] -> ISPCCompilerM [BlockItem]
compileGetStructVals Name
fstruct [VName]
free_args [(Type, ValueType)]
free_ctypes
            body' <- CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
GC.collect (CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem])
-> CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$ MCCode -> CompilerM Multicore ISPCState ()
compileCode MCCode
body
            mapM_ GC.item decl_cached
            mapM_ GC.item =<< GC.declAllocatedMem

            -- Make inner kernel for error handling, if needed
            if mayProduceError body
              then do
                fatmems <- gets (map fst . GC.compDeclaredMem)
                mstruct <- prepareMemStruct lexmems fatmems
                let mstruct' = Name
mstruct Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_"
                innerShim <- ispcDef "inner_ispc" $ \Name
t -> do
                  innerBody <- CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
GC.collect (CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem])
-> CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$ do
                    InitGroup -> CompilerM Multicore ISPCState ()
forall op s. InitGroup -> CompilerM op s ()
GC.decl [C.cdecl|$tyqual:uniform struct futhark_context * $tyqual:uniform ctx = $id:fstruct'->ctx;|]
                    [BlockItem] -> CompilerM Multicore ISPCState ()
forall op s. [BlockItem] -> CompilerM op s ()
GC.items ([BlockItem] -> CompilerM Multicore ISPCState ())
-> ISPCCompilerM [BlockItem] -> CompilerM Multicore ISPCState ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> [VName] -> [(Type, ValueType)] -> ISPCCompilerM [BlockItem]
compileGetStructVals Name
fstruct [VName]
free_args [(Type, ValueType)]
free_ctypes
                    Name
-> [(VName, VName)] -> [VName] -> CompilerM Multicore ISPCState ()
compileGetMemStructVals Name
mstruct' [(VName, VName)]
lexmems [VName]
fatmems
                    InitGroup -> CompilerM Multicore ISPCState ()
forall op s. InitGroup -> CompilerM op s ()
GC.decl [C.cdecl|$tyqual:uniform int err = 0;|]
                    (BlockItem -> CompilerM Multicore ISPCState ())
-> [BlockItem] -> CompilerM Multicore ISPCState ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ BlockItem -> CompilerM Multicore ISPCState ()
forall op s. BlockItem -> CompilerM op s ()
GC.item [BlockItem]
body'
                    Name
-> [(VName, VName)] -> [VName] -> CompilerM Multicore ISPCState ()
compileWritebackMemStructVals Name
mstruct' [(VName, VName)]
lexmems [VName]
fatmems
                    Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|return err;|]
                  pure
                    [C.cedecl|
                static $tyqual:unmasked inline $tyqual:uniform int $id:t(
                  $tyqual:uniform typename int64_t start,
                  $tyqual:uniform typename int64_t end,
                  struct $id:fstruct $tyqual:uniform * $tyqual:uniform $id:fstruct',
                  struct $id:mstruct $tyqual:uniform * $tyqual:uniform $id:mstruct') {
                  $items:innerBody
                }|]
                -- Call the kernel and read back potentially changed memory
                GC.decl [C.cdecl|$tyqual:uniform int err = $id:innerShim(start, end, $id:fstruct', &$id:mstruct');|]
                compileReadbackMemStructVals mstruct' lexmems fatmems
              else do
                GC.decl [C.cdecl|$tyqual:uniform int err = 0;|]
                mapM_ GC.item body'

            free_mem <- freeAllocatedMem
            GC.stm [C.cstm|cleanup: {$stms:free_cached $items:free_mem}|]
            GC.stm [C.cstm|return err;|]
    GC.earlyDecl
      [C.cedecl|int $id:s(typename int64_t start,
                                  typename int64_t end,
                                  struct $id:fstruct * $id:fstruct');|]
    pure
      [C.cedecl|
        $tyqual:export $tyqual:uniform int $id:s($tyqual:uniform typename int64_t start,
                                                 $tyqual:uniform typename int64_t end,
                                                 struct $id:fstruct $tyqual:uniform * $tyqual:uniform $id:fstruct' ) {
          $items:mainBody
        }|]

  -- Generate C code to call into ISPC kernel
  GC.items
    [C.citems|
    err = $id:ispcShim(start, end, & $id:fstruct');
    if (err != 0) {
      goto cleanup;
    }|]
compileOp (ForEach VName
i Exp
from Exp
bound MCCode
body) = do
  from' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
from
  bound' <- compileExp bound
  body' <- GC.collect $ compileCode body
  if mayProduceError body
    then
      GC.stms
        [C.cstms|
      for ($tyqual:uniform typename int64_t i = 0; i < (($exp:bound' - $exp:from') / programCount); i++) {
        typename int64_t $id:i = $exp:from' + programIndex + i * programCount;
        $items:body'
      }
      if (programIndex < (($exp:bound' - $exp:from') % programCount)) {
        typename int64_t $id:i = $exp:from' + programIndex + ((($exp:bound' - $exp:from') / programCount) * programCount);
        $items:body'
      }|]
    else
      GC.stms
        [C.cstms|
      $escstm:(T.unpack ("foreach (" <> prettyText i <> " = " <> expText from' <> " ... " <> expText bound' <> ")")) {
        $items:body'
      }|]
compileOp (ForEachActive VName
name MCCode
body) = do
  body' <- CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
GC.collect (CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem])
-> CompilerM Multicore ISPCState () -> ISPCCompilerM [BlockItem]
forall a b. (a -> b) -> a -> b
$ MCCode -> CompilerM Multicore ISPCState ()
compileCode MCCode
body
  GC.stms
    [C.cstms|
    for ($tyqual:uniform unsigned int $id:name = 0; $id:name < programCount; $id:name++) {
      if (programIndex == $id:name) {
        $items:body'
      }
    }|]
compileOp (ExtractLane VName
dest (ValueExp PrimValue
v) Exp
_) =
  -- extract() on constants is not allowed (type is uniform, not
  -- varying), so just turn them into an assignment.
  Stm -> CompilerM Multicore ISPCState ()
forall op s. Stm -> CompilerM op s ()
GC.stm [C.cstm|$id:dest = $exp:v;|]
compileOp (ExtractLane VName
dest Exp
tar Exp
lane) = do
  tar' <- Exp -> CompilerM Multicore ISPCState Exp
compileExp Exp
tar
  lane' <- compileExp lane
  GC.stm [C.cstm|$id:dest = extract($exp:tar', $exp:lane');|]
compileOp (Atomic AtomicOp
aop) =
  AtomicOp
-> (Type -> VName -> ISPCCompilerM Type)
-> CompilerM Multicore ISPCState ()
forall op s.
AtomicOp
-> (Type -> VName -> CompilerM op s Type) -> CompilerM op s ()
MC.atomicOps AtomicOp
aop ((Type -> VName -> ISPCCompilerM Type)
 -> CompilerM Multicore ISPCState ())
-> (Type -> VName -> ISPCCompilerM Type)
-> CompilerM Multicore ISPCState ()
forall a b. (a -> b) -> a -> b
$ \Type
ty VName
arr -> do
    cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM Multicore ISPCState (Maybe VName)
-> CompilerM Multicore ISPCState Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM Multicore ISPCState (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
GC.cacheMem VName
arr
    if cached
      then pure [C.cty|$tyqual:varying $ty:ty* $tyqual:uniform|]
      else pure [C.cty|$ty:ty*|]
compileOp Multicore
op = OpCompiler Multicore ISPCState
forall s. OpCompiler Multicore s
MC.compileOp Multicore
op

-- | Like @GenericC.cachingMemory@, but adapted for ISPC codegen.
cachingMemory ::
  M.Map VName Space ->
  ([C.BlockItem] -> [C.Stm] -> [(VName, VName)] -> GC.CompilerM op s a) ->
  GC.CompilerM op s a
cachingMemory :: forall op s a.
Map VName Space
-> ([BlockItem] -> [Stm] -> [(VName, VName)] -> CompilerM op s a)
-> CompilerM op s a
cachingMemory Map VName Space
lexical [BlockItem] -> [Stm] -> [(VName, VName)] -> CompilerM op s a
f = do
  let cached :: [VName]
cached = Map VName Space -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName Space -> [VName]) -> Map VName Space -> [VName]
forall a b. (a -> b) -> a -> b
$ (Space -> Bool) -> Map VName Space -> Map VName Space
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Space -> Space -> Bool
forall a. Eq a => a -> a -> Bool
== Space
DefaultSpace) Map VName Space
lexical

  cached' <- [VName]
-> (VName -> CompilerM op s (VName, VName))
-> CompilerM op s [(VName, VName)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
cached ((VName -> CompilerM op s (VName, VName))
 -> CompilerM op s [(VName, VName)])
-> (VName -> CompilerM op s (VName, VName))
-> CompilerM op s [(VName, VName)]
forall a b. (a -> b) -> a -> b
$ \VName
mem -> do
    size <- String -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> CompilerM op s VName) -> String -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ VName -> String
forall a. Pretty a => a -> String
prettyString VName
mem String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_cached_size"
    pure (mem, size)

  let lexMem CompilerEnv op s
env =
        CompilerEnv op s
env
          { GC.envCachedMem =
              M.fromList (map (first (`C.toExp` noLoc)) cached')
                <> GC.envCachedMem env
          }

      declCached (a
mem, a
size) =
        [ [C.citem|size_t $id:size = 0;|],
          [C.citem|$tyqual:varying unsigned char * $tyqual:uniform $id:mem = NULL;|]
        ]

      freeCached (a
mem, b
_) =
        [C.cstm|free($id:mem);|]

  local lexMem $ f (concatMap declCached cached') (map freeCached cached') cached'

-- Variability analysis
type Dependencies = M.Map VName Names

data Variability = Uniform | Varying
  deriving (Variability -> Variability -> Bool
(Variability -> Variability -> Bool)
-> (Variability -> Variability -> Bool) -> Eq Variability
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Variability -> Variability -> Bool
== :: Variability -> Variability -> Bool
$c/= :: Variability -> Variability -> Bool
/= :: Variability -> Variability -> Bool
Eq, Eq Variability
Eq Variability =>
(Variability -> Variability -> Ordering)
-> (Variability -> Variability -> Bool)
-> (Variability -> Variability -> Bool)
-> (Variability -> Variability -> Bool)
-> (Variability -> Variability -> Bool)
-> (Variability -> Variability -> Variability)
-> (Variability -> Variability -> Variability)
-> Ord Variability
Variability -> Variability -> Bool
Variability -> Variability -> Ordering
Variability -> Variability -> Variability
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Variability -> Variability -> Ordering
compare :: Variability -> Variability -> Ordering
$c< :: Variability -> Variability -> Bool
< :: Variability -> Variability -> Bool
$c<= :: Variability -> Variability -> Bool
<= :: Variability -> Variability -> Bool
$c> :: Variability -> Variability -> Bool
> :: Variability -> Variability -> Bool
$c>= :: Variability -> Variability -> Bool
>= :: Variability -> Variability -> Bool
$cmax :: Variability -> Variability -> Variability
max :: Variability -> Variability -> Variability
$cmin :: Variability -> Variability -> Variability
min :: Variability -> Variability -> Variability
Ord, Int -> Variability -> String -> String
[Variability] -> String -> String
Variability -> String
(Int -> Variability -> String -> String)
-> (Variability -> String)
-> ([Variability] -> String -> String)
-> Show Variability
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Variability -> String -> String
showsPrec :: Int -> Variability -> String -> String
$cshow :: Variability -> String
show :: Variability -> String
$cshowList :: [Variability] -> String -> String
showList :: [Variability] -> String -> String
Show)

newtype VariabilityM a
  = VariabilityM (ReaderT Names (State Dependencies) a)
  deriving
    ( (forall a b. (a -> b) -> VariabilityM a -> VariabilityM b)
-> (forall a b. a -> VariabilityM b -> VariabilityM a)
-> Functor VariabilityM
forall a b. a -> VariabilityM b -> VariabilityM a
forall a b. (a -> b) -> VariabilityM a -> VariabilityM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> VariabilityM a -> VariabilityM b
fmap :: forall a b. (a -> b) -> VariabilityM a -> VariabilityM b
$c<$ :: forall a b. a -> VariabilityM b -> VariabilityM a
<$ :: forall a b. a -> VariabilityM b -> VariabilityM a
Functor,
      Functor VariabilityM
Functor VariabilityM =>
(forall a. a -> VariabilityM a)
-> (forall a b.
    VariabilityM (a -> b) -> VariabilityM a -> VariabilityM b)
-> (forall a b c.
    (a -> b -> c)
    -> VariabilityM a -> VariabilityM b -> VariabilityM c)
-> (forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b)
-> (forall a b. VariabilityM a -> VariabilityM b -> VariabilityM a)
-> Applicative VariabilityM
forall a. a -> VariabilityM a
forall a b. VariabilityM a -> VariabilityM b -> VariabilityM a
forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b
forall a b.
VariabilityM (a -> b) -> VariabilityM a -> VariabilityM b
forall a b c.
(a -> b -> c) -> VariabilityM a -> VariabilityM b -> VariabilityM c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> VariabilityM a
pure :: forall a. a -> VariabilityM a
$c<*> :: forall a b.
VariabilityM (a -> b) -> VariabilityM a -> VariabilityM b
<*> :: forall a b.
VariabilityM (a -> b) -> VariabilityM a -> VariabilityM b
$cliftA2 :: forall a b c.
(a -> b -> c) -> VariabilityM a -> VariabilityM b -> VariabilityM c
liftA2 :: forall a b c.
(a -> b -> c) -> VariabilityM a -> VariabilityM b -> VariabilityM c
$c*> :: forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b
*> :: forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b
$c<* :: forall a b. VariabilityM a -> VariabilityM b -> VariabilityM a
<* :: forall a b. VariabilityM a -> VariabilityM b -> VariabilityM a
Applicative,
      Applicative VariabilityM
Applicative VariabilityM =>
(forall a b.
 VariabilityM a -> (a -> VariabilityM b) -> VariabilityM b)
-> (forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b)
-> (forall a. a -> VariabilityM a)
-> Monad VariabilityM
forall a. a -> VariabilityM a
forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b
forall a b.
VariabilityM a -> (a -> VariabilityM b) -> VariabilityM b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b.
VariabilityM a -> (a -> VariabilityM b) -> VariabilityM b
>>= :: forall a b.
VariabilityM a -> (a -> VariabilityM b) -> VariabilityM b
$c>> :: forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b
>> :: forall a b. VariabilityM a -> VariabilityM b -> VariabilityM b
$creturn :: forall a. a -> VariabilityM a
return :: forall a. a -> VariabilityM a
Monad,
      MonadState Dependencies,
      MonadReader Names
    )

execVariabilityM :: VariabilityM a -> Dependencies
execVariabilityM :: forall a. VariabilityM a -> Dependencies
execVariabilityM (VariabilityM ReaderT Names (State Dependencies) a
m) =
  State Dependencies a -> Dependencies -> Dependencies
forall s a. State s a -> s -> s
execState (ReaderT Names (State Dependencies) a
-> Names -> State Dependencies a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Names (State Dependencies) a
m Names
forall a. Monoid a => a
mempty) Dependencies
forall a. Monoid a => a
mempty

-- | Extend the set of dependencies with a new one
addDeps :: VName -> Names -> VariabilityM ()
addDeps :: VName -> Names -> VariabilityM ()
addDeps VName
v Names
ns = do
  deps <- VariabilityM Dependencies
forall s (m :: * -> *). MonadState s m => m s
get
  env <- ask
  case M.lookup v deps of
    Maybe Names
Nothing -> Dependencies -> VariabilityM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Dependencies -> VariabilityM ())
-> Dependencies -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ VName -> Names -> Dependencies -> Dependencies
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (Names
ns Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
env) Dependencies
deps
    Just Names
ns' -> Dependencies -> VariabilityM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Dependencies -> VariabilityM ())
-> Dependencies -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ VName -> Names -> Dependencies -> Dependencies
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (Names
ns Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
ns') Dependencies
deps

-- | Find all the dependencies in a body of code
findDeps :: MCCode -> VariabilityM ()
findDeps :: MCCode -> VariabilityM ()
findDeps (MCCode
x :>>: MCCode
y) = do
  MCCode -> VariabilityM ()
findDeps MCCode
x
  MCCode -> VariabilityM ()
findDeps MCCode
y
findDeps (If TExp Bool
cond MCCode
x MCCode
y) =
  (Names -> Names) -> VariabilityM () -> VariabilityM ()
forall a. (Names -> Names) -> VariabilityM a -> VariabilityM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> TExp Bool -> Names
forall a. FreeIn a => a -> Names
freeIn TExp Bool
cond) (VariabilityM () -> VariabilityM ())
-> VariabilityM () -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ do
    MCCode -> VariabilityM ()
findDeps MCCode
x
    MCCode -> VariabilityM ()
findDeps MCCode
y
findDeps (For VName
idx Exp
bound MCCode
x) = do
  VName -> Names -> VariabilityM ()
addDeps VName
idx Names
free
  (Names -> Names) -> VariabilityM () -> VariabilityM ()
forall a. (Names -> Names) -> VariabilityM a -> VariabilityM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
free) (VariabilityM () -> VariabilityM ())
-> VariabilityM () -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ MCCode -> VariabilityM ()
findDeps MCCode
x
  where
    free :: Names
free = Exp -> Names
forall a. FreeIn a => a -> Names
freeIn Exp
bound
findDeps (While TExp Bool
cond MCCode
x) = do
  (Names -> Names) -> VariabilityM () -> VariabilityM ()
forall a. (Names -> Names) -> VariabilityM a -> VariabilityM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> TExp Bool -> Names
forall a. FreeIn a => a -> Names
freeIn TExp Bool
cond) (VariabilityM () -> VariabilityM ())
-> VariabilityM () -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ MCCode -> VariabilityM ()
findDeps MCCode
x
findDeps (Op (SegOp String
_ [Param]
free ParallelTask
_ Maybe ParallelTask
_ [Param]
retvals SchedulerInfo
_)) =
  (Param -> VariabilityM ()) -> [Param] -> VariabilityM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_
    ( \Param
x ->
        VName -> Names -> VariabilityM ()
addDeps (Param -> VName
paramName Param
x) (Names -> VariabilityM ()) -> Names -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$
          [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$
            (Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
free
    )
    [Param]
retvals
findDeps (Op (ForEach VName
_ Exp
_ Exp
_ MCCode
body)) =
  MCCode -> VariabilityM ()
findDeps MCCode
body
findDeps (Op (ForEachActive VName
_ MCCode
body)) =
  MCCode -> VariabilityM ()
findDeps MCCode
body
findDeps (SetScalar VName
name Exp
e) =
  VName -> Names -> VariabilityM ()
addDeps VName
name (Names -> VariabilityM ()) -> Names -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ Exp -> Names
forall a. FreeIn a => a -> Names
freeIn Exp
e
findDeps (Call [VName]
tars Name
_ [Arg]
args) =
  (VName -> VariabilityM ()) -> [VName] -> VariabilityM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\VName
x -> VName -> Names -> VariabilityM ()
addDeps VName
x (Names -> VariabilityM ()) -> Names -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ [Arg] -> Names
forall a. FreeIn a => a -> Names
freeIn [Arg]
args) [VName]
tars
findDeps (Read VName
x VName
arr (Count TExp Int64
iexp) PrimType
_ Space
DefaultSpace Volatility
_) = do
  VName -> Names -> VariabilityM ()
addDeps VName
x (Names -> VariabilityM ()) -> Names -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ Exp -> Names
forall a. FreeIn a => a -> Names
freeIn (TExp Int64 -> Exp
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TExp Int64
iexp)
  VName -> Names -> VariabilityM ()
addDeps VName
x (Names -> VariabilityM ()) -> Names -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
arr
findDeps (Op (GetLoopBounds VName
x VName
y)) = do
  VName -> Names -> VariabilityM ()
addDeps VName
x Names
forall a. Monoid a => a
mempty
  VName -> Names -> VariabilityM ()
addDeps VName
y Names
forall a. Monoid a => a
mempty
findDeps (Op (ExtractLane VName
x Exp
_ Exp
_)) = do
  VName -> Names -> VariabilityM ()
addDeps VName
x Names
forall a. Monoid a => a
mempty
findDeps (Op (Atomic (AtomicCmpXchg PrimType
_ VName
old VName
arr Count Elements (TExp Int32)
ind VName
res Exp
val))) = do
  VName -> Names -> VariabilityM ()
addDeps VName
res (Names -> VariabilityM ()) -> Names -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ VName -> Names
forall a. FreeIn a => a -> Names
freeIn VName
arr Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> Names
forall a. FreeIn a => a -> Names
freeIn Count Elements (TExp Int32)
ind Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Exp -> Names
forall a. FreeIn a => a -> Names
freeIn Exp
val
  VName -> Names -> VariabilityM ()
addDeps VName
old (Names -> VariabilityM ()) -> Names -> VariabilityM ()
forall a b. (a -> b) -> a -> b
$ VName -> Names
forall a. FreeIn a => a -> Names
freeIn VName
arr Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> Names
forall a. FreeIn a => a -> Names
freeIn Count Elements (TExp Int32)
ind Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Exp -> Names
forall a. FreeIn a => a -> Names
freeIn Exp
val
findDeps MCCode
_ = () -> VariabilityM ()
forall a. a -> VariabilityM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Take a list of dependencies and iterate them to a fixed point.
depsFixedPoint :: Dependencies -> Dependencies
depsFixedPoint :: Dependencies -> Dependencies
depsFixedPoint Dependencies
deps =
  if Dependencies
deps Dependencies -> Dependencies -> Bool
forall a. Eq a => a -> a -> Bool
== Dependencies
deps'
    then Dependencies
deps
    else Dependencies -> Dependencies
depsFixedPoint Dependencies
deps'
  where
    grow :: Names -> Names
grow Names
names =
      Names
names Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> (VName -> Names) -> IntMap VName -> Names
forall m a. Monoid m => (a -> m) -> IntMap a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\VName
n -> Names -> VName -> Dependencies -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
n Dependencies
deps) (Names -> IntMap VName
namesIntMap Names
names)
    deps' :: Dependencies
deps' = (Names -> Names) -> Dependencies -> Dependencies
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Names -> Names
grow Dependencies
deps

-- | Find roots of variance. These are memory blocks declared in
-- the current scope as well as loop indices of foreach loops.
findVarying :: MCCode -> [VName]
findVarying :: MCCode -> [VName]
findVarying (MCCode
x :>>: MCCode
y) = MCCode -> [VName]
findVarying MCCode
x [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ MCCode -> [VName]
findVarying MCCode
y
findVarying (If TExp Bool
_ MCCode
x MCCode
y) = MCCode -> [VName]
findVarying MCCode
x [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ MCCode -> [VName]
findVarying MCCode
y
findVarying (For VName
_ Exp
_ MCCode
x) = MCCode -> [VName]
findVarying MCCode
x
findVarying (While TExp Bool
_ MCCode
x) = MCCode -> [VName]
findVarying MCCode
x
findVarying (Op (ForEachActive VName
_ MCCode
body)) = MCCode -> [VName]
findVarying MCCode
body
findVarying (Op (ForEach VName
idx Exp
_ Exp
_ MCCode
body)) = VName
idx VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: MCCode -> [VName]
findVarying MCCode
body
findVarying (DeclareMem VName
mem Space
_) = [VName
mem]
findVarying MCCode
_ = []

-- | Analyze variability in a body of code and run an action with
-- info about that variability in the compiler state.
analyzeVariability :: MCCode -> ISPCCompilerM a -> ISPCCompilerM a
analyzeVariability :: forall a. MCCode -> ISPCCompilerM a -> ISPCCompilerM a
analyzeVariability MCCode
code ISPCCompilerM a
m = do
  let roots :: [VName]
roots = MCCode -> [VName]
findVarying MCCode
code
  let deps :: Dependencies
deps = Dependencies -> Dependencies
depsFixedPoint (Dependencies -> Dependencies) -> Dependencies -> Dependencies
forall a b. (a -> b) -> a -> b
$ VariabilityM () -> Dependencies
forall a. VariabilityM a -> Dependencies
execVariabilityM (VariabilityM () -> Dependencies)
-> VariabilityM () -> Dependencies
forall a b. (a -> b) -> a -> b
$ MCCode -> VariabilityM ()
findDeps MCCode
code
  let safelist :: Dependencies
safelist = (Names -> Bool) -> Dependencies -> Dependencies
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (\Names
b -> (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Names -> Bool
`notNameIn` Names
b) [VName]
roots) Dependencies
deps
  let safe :: Names
safe = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Dependencies -> [VName]
forall k a. Map k a -> [k]
M.keys Dependencies
safelist
  pre_state <- CompilerM Multicore ISPCState ISPCState
forall op s. CompilerM op s s
GC.getUserState
  GC.modifyUserState (\ISPCState
s -> ISPCState
s {sUniform = safe})
  a <- m
  GC.modifyUserState (\ISPCState
s -> ISPCState
s {sUniform = sUniform pre_state})
  pure a

-- | Get the variability of a variable
getVariability :: VName -> ISPCCompilerM Variability
getVariability :: VName -> ISPCCompilerM Variability
getVariability VName
name = do
  uniforms <- ISPCState -> Names
sUniform (ISPCState -> Names)
-> CompilerM Multicore ISPCState ISPCState
-> CompilerM Multicore ISPCState Names
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CompilerM Multicore ISPCState ISPCState
forall op s. CompilerM op s s
GC.getUserState
  pure $
    if name `nameIn` uniforms
      then Uniform
      else Varying

-- | Get the variability qualifiers of a variable
getVariabilityQuals :: VName -> ISPCCompilerM [C.TypeQual]
getVariabilityQuals :: VName -> ISPCCompilerM [TypeQual]
getVariabilityQuals VName
name = Variability -> [TypeQual]
variQuals (Variability -> [TypeQual])
-> ISPCCompilerM Variability -> ISPCCompilerM [TypeQual]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ISPCCompilerM Variability
getVariability VName
name
  where
    variQuals :: Variability -> [TypeQual]
variQuals Variability
Uniform = [C.ctyquals|$tyqual:uniform|]
    variQuals Variability
Varying = []