{-# LANGUAGE QuasiQuotes #-}

-- | C code generator framework.
module Futhark.CodeGen.Backends.GenericC.Monad
  ( -- * Pluggable compiler
    Operations (..),
    Publicness (..),
    OpCompiler,
    ErrorCompiler,
    CallCompiler,
    PointerQuals,
    MemoryType,
    WriteScalar,
    writeScalarPointerWithQuals,
    ReadScalar,
    readScalarPointerWithQuals,
    Allocate,
    Deallocate,
    CopyBarrier (..),
    Copy,
    DoCopy,

    -- * Monadic compiler interface
    CompilerM,
    CompilerState (..),
    CompilerEnv (..),
    getUserState,
    modifyUserState,
    generateProgramStruct,
    runCompilerM,
    inNewFunction,
    cachingMemory,
    volQuals,
    rawMem,
    item,
    items,
    stm,
    stms,
    decl,
    headerDecl,
    publicDef,
    publicDef_,
    onClear,
    HeaderSection (..),
    libDecl,
    earlyDecl,
    publicName,
    contextField,
    contextFieldDyn,
    memToCType,
    cacheMem,
    fatMemory,
    rawMemCType,
    freeRawMem,
    allocRawMem,
    fatMemType,
    declAllocatedMem,
    freeAllocatedMem,
    collect,
    collect',
    contextType,
    configType,

    -- * Building Blocks
    copyMemoryDefaultSpace,
    derefPointer,
    setMem,
    allocMem,
    unRefMem,
    declMem,
    resetMem,
    fatMemAlloc,
    fatMemSet,
    fatMemUnRef,
    criticalSection,
    module Futhark.CodeGen.Backends.SimpleRep,
  )
where

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor (first)
import Data.DList qualified as DL
import Data.List (unzip4)
import Data.Loc
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericC.Pretty
import Futhark.CodeGen.Backends.SimpleRep
import Futhark.CodeGen.ImpCode
import Futhark.MonadFreshNames
import Language.C.Quote.OpenCL qualified as C
import Language.C.Syntax qualified as C

-- How public an array type definition sould be.  Public types show up
-- in the generated API, while private types are used only to
-- implement the members of opaques.
data Publicness = Private | Public
  deriving (Publicness -> Publicness -> Bool
(Publicness -> Publicness -> Bool)
-> (Publicness -> Publicness -> Bool) -> Eq Publicness
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Publicness -> Publicness -> Bool
== :: Publicness -> Publicness -> Bool
$c/= :: Publicness -> Publicness -> Bool
/= :: Publicness -> Publicness -> Bool
Eq, Eq Publicness
Eq Publicness =>
(Publicness -> Publicness -> Ordering)
-> (Publicness -> Publicness -> Bool)
-> (Publicness -> Publicness -> Bool)
-> (Publicness -> Publicness -> Bool)
-> (Publicness -> Publicness -> Bool)
-> (Publicness -> Publicness -> Publicness)
-> (Publicness -> Publicness -> Publicness)
-> Ord Publicness
Publicness -> Publicness -> Bool
Publicness -> Publicness -> Ordering
Publicness -> Publicness -> Publicness
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 :: Publicness -> Publicness -> Ordering
compare :: Publicness -> Publicness -> Ordering
$c< :: Publicness -> Publicness -> Bool
< :: Publicness -> Publicness -> Bool
$c<= :: Publicness -> Publicness -> Bool
<= :: Publicness -> Publicness -> Bool
$c> :: Publicness -> Publicness -> Bool
> :: Publicness -> Publicness -> Bool
$c>= :: Publicness -> Publicness -> Bool
>= :: Publicness -> Publicness -> Bool
$cmax :: Publicness -> Publicness -> Publicness
max :: Publicness -> Publicness -> Publicness
$cmin :: Publicness -> Publicness -> Publicness
min :: Publicness -> Publicness -> Publicness
Ord, Int -> Publicness -> ShowS
[Publicness] -> ShowS
Publicness -> [Char]
(Int -> Publicness -> ShowS)
-> (Publicness -> [Char])
-> ([Publicness] -> ShowS)
-> Show Publicness
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Publicness -> ShowS
showsPrec :: Int -> Publicness -> ShowS
$cshow :: Publicness -> [Char]
show :: Publicness -> [Char]
$cshowList :: [Publicness] -> ShowS
showList :: [Publicness] -> ShowS
Show)

type ArrayType = (Signedness, PrimType, Int)

data CompilerState s = CompilerState
  { forall s. CompilerState s -> Map ArrayType Publicness
compArrayTypes :: M.Map ArrayType Publicness,
    forall s. CompilerState s -> DList Definition
compEarlyDecls :: DL.DList C.Definition,
    forall s. CompilerState s -> VNameSource
compNameSrc :: VNameSource,
    forall s. CompilerState s -> s
compUserState :: s,
    forall s. CompilerState s -> Map HeaderSection (DList Definition)
compHeaderDecls :: M.Map HeaderSection (DL.DList C.Definition),
    forall s. CompilerState s -> DList Definition
compLibDecls :: DL.DList C.Definition,
    forall s.
CompilerState s -> DList (Id, Type, Maybe Exp, Maybe (Stm, Stm))
compCtxFields :: DL.DList (C.Id, C.Type, Maybe C.Exp, Maybe (C.Stm, C.Stm)),
    forall s. CompilerState s -> DList BlockItem
compClearItems :: DL.DList C.BlockItem,
    forall s. CompilerState s -> [(VName, Space)]
compDeclaredMem :: [(VName, Space)],
    forall s. CompilerState s -> DList BlockItem
compItems :: DL.DList C.BlockItem
  }

newCompilerState :: VNameSource -> s -> CompilerState s
newCompilerState :: forall s. VNameSource -> s -> CompilerState s
newCompilerState VNameSource
src s
s =
  CompilerState
    { compArrayTypes :: Map ArrayType Publicness
compArrayTypes = Map ArrayType Publicness
forall a. Monoid a => a
mempty,
      compEarlyDecls :: DList Definition
compEarlyDecls = DList Definition
forall a. Monoid a => a
mempty,
      compNameSrc :: VNameSource
compNameSrc = VNameSource
src,
      compUserState :: s
compUserState = s
s,
      compHeaderDecls :: Map HeaderSection (DList Definition)
compHeaderDecls = Map HeaderSection (DList Definition)
forall a. Monoid a => a
mempty,
      compLibDecls :: DList Definition
compLibDecls = DList Definition
forall a. Monoid a => a
mempty,
      compCtxFields :: DList (Id, Type, Maybe Exp, Maybe (Stm, Stm))
compCtxFields = DList (Id, Type, Maybe Exp, Maybe (Stm, Stm))
forall a. Monoid a => a
mempty,
      compClearItems :: DList BlockItem
compClearItems = DList BlockItem
forall a. Monoid a => a
mempty,
      compDeclaredMem :: [(VName, Space)]
compDeclaredMem = [(VName, Space)]
forall a. Monoid a => a
mempty,
      compItems :: DList BlockItem
compItems = DList BlockItem
forall a. Monoid a => a
mempty
    }

-- | In which part of the header file we put the declaration.  This is
-- to ensure that the header file remains structured and readable.
data HeaderSection
  = ArrayDecl Name
  | OpaqueTypeDecl Name
  | OpaqueDecl Name
  | EntryDecl
  | MiscDecl
  | InitDecl
  deriving (HeaderSection -> HeaderSection -> Bool
(HeaderSection -> HeaderSection -> Bool)
-> (HeaderSection -> HeaderSection -> Bool) -> Eq HeaderSection
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: HeaderSection -> HeaderSection -> Bool
== :: HeaderSection -> HeaderSection -> Bool
$c/= :: HeaderSection -> HeaderSection -> Bool
/= :: HeaderSection -> HeaderSection -> Bool
Eq, Eq HeaderSection
Eq HeaderSection =>
(HeaderSection -> HeaderSection -> Ordering)
-> (HeaderSection -> HeaderSection -> Bool)
-> (HeaderSection -> HeaderSection -> Bool)
-> (HeaderSection -> HeaderSection -> Bool)
-> (HeaderSection -> HeaderSection -> Bool)
-> (HeaderSection -> HeaderSection -> HeaderSection)
-> (HeaderSection -> HeaderSection -> HeaderSection)
-> Ord HeaderSection
HeaderSection -> HeaderSection -> Bool
HeaderSection -> HeaderSection -> Ordering
HeaderSection -> HeaderSection -> HeaderSection
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 :: HeaderSection -> HeaderSection -> Ordering
compare :: HeaderSection -> HeaderSection -> Ordering
$c< :: HeaderSection -> HeaderSection -> Bool
< :: HeaderSection -> HeaderSection -> Bool
$c<= :: HeaderSection -> HeaderSection -> Bool
<= :: HeaderSection -> HeaderSection -> Bool
$c> :: HeaderSection -> HeaderSection -> Bool
> :: HeaderSection -> HeaderSection -> Bool
$c>= :: HeaderSection -> HeaderSection -> Bool
>= :: HeaderSection -> HeaderSection -> Bool
$cmax :: HeaderSection -> HeaderSection -> HeaderSection
max :: HeaderSection -> HeaderSection -> HeaderSection
$cmin :: HeaderSection -> HeaderSection -> HeaderSection
min :: HeaderSection -> HeaderSection -> HeaderSection
Ord)

-- | A substitute expression compiler, tried before the main
-- compilation function.
type OpCompiler op s = op -> CompilerM op s ()

type ErrorCompiler op s = ErrorMsg Exp -> String -> CompilerM op s ()

-- | The address space qualifiers for a pointer of the given type with
-- the given annotation.
type PointerQuals = String -> [C.TypeQual]

-- | The type of a memory block in the given memory space.
type MemoryType op s = SpaceId -> CompilerM op s C.Type

-- | Write a scalar to the given memory block with the given element
-- index and in the given memory space.
type WriteScalar op s =
  C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s ()

-- | Read a scalar from the given memory block with the given element
-- index and in the given memory space.
type ReadScalar op s =
  C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp

-- | Allocate a memory block of the given size and with the given tag
-- in the given memory space, saving a reference in the given variable
-- name.
type Allocate op s =
  C.Exp ->
  C.Exp ->
  C.Exp ->
  SpaceId ->
  CompilerM op s ()

-- | De-allocate the given memory block, with the given tag, with the
-- given size,, which is in the given memory space.
type Deallocate op s = C.Exp -> C.Exp -> C.Exp -> SpaceId -> CompilerM op s ()

-- | Whether a copying operation should implicitly function as a
-- barrier regarding further operations on the source.  This is a
-- rather subtle detail and is mostly useful for letting some
-- device/GPU copies be asynchronous (#1664).
data CopyBarrier
  = CopyBarrier
  | -- | Explicit context synchronisation should be done
    -- before the source or target is used.
    CopyNoBarrier
  deriving (CopyBarrier -> CopyBarrier -> Bool
(CopyBarrier -> CopyBarrier -> Bool)
-> (CopyBarrier -> CopyBarrier -> Bool) -> Eq CopyBarrier
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CopyBarrier -> CopyBarrier -> Bool
== :: CopyBarrier -> CopyBarrier -> Bool
$c/= :: CopyBarrier -> CopyBarrier -> Bool
/= :: CopyBarrier -> CopyBarrier -> Bool
Eq, Int -> CopyBarrier -> ShowS
[CopyBarrier] -> ShowS
CopyBarrier -> [Char]
(Int -> CopyBarrier -> ShowS)
-> (CopyBarrier -> [Char])
-> ([CopyBarrier] -> ShowS)
-> Show CopyBarrier
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CopyBarrier -> ShowS
showsPrec :: Int -> CopyBarrier -> ShowS
$cshow :: CopyBarrier -> [Char]
show :: CopyBarrier -> [Char]
$cshowList :: [CopyBarrier] -> ShowS
showList :: [CopyBarrier] -> ShowS
Show)

-- | Copy from one memory block to another.
type Copy op s =
  CopyBarrier ->
  C.Exp ->
  C.Exp ->
  Space ->
  C.Exp ->
  C.Exp ->
  Space ->
  C.Exp ->
  CompilerM op s ()

-- | Perform an 'Copy'.  It is expected that these functions are
-- each specialised on which spaces they operate on, so that is not part of their arguments.
type DoCopy op s =
  CopyBarrier ->
  PrimType ->
  [Count Elements C.Exp] ->
  C.Exp ->
  ( Count Elements C.Exp,
    [Count Elements C.Exp]
  ) ->
  C.Exp ->
  ( Count Elements C.Exp,
    [Count Elements C.Exp]
  ) ->
  CompilerM op s ()

-- | Call a function.
type CallCompiler op s = [VName] -> Name -> [C.Exp] -> CompilerM op s ()

data Operations op s = Operations
  { forall op s. Operations op s -> WriteScalar op s
opsWriteScalar :: WriteScalar op s,
    forall op s. Operations op s -> ReadScalar op s
opsReadScalar :: ReadScalar op s,
    forall op s. Operations op s -> Allocate op s
opsAllocate :: Allocate op s,
    forall op s. Operations op s -> Allocate op s
opsDeallocate :: Deallocate op s,
    forall op s. Operations op s -> Copy op s
opsCopy :: Copy op s,
    forall op s. Operations op s -> MemoryType op s
opsMemoryType :: MemoryType op s,
    forall op s. Operations op s -> OpCompiler op s
opsCompiler :: OpCompiler op s,
    forall op s. Operations op s -> ErrorCompiler op s
opsError :: ErrorCompiler op s,
    forall op s. Operations op s -> CallCompiler op s
opsCall :: CallCompiler op s,
    -- | @(dst,src)@-space mapping to copy functions.
    forall op s. Operations op s -> Map (Space, Space) (DoCopy op s)
opsCopies :: M.Map (Space, Space) (DoCopy op s),
    -- | If true, use reference counting.  Otherwise, bare
    -- pointers.
    forall op s. Operations op s -> Bool
opsFatMemory :: Bool,
    -- | Code to bracket critical sections.
    forall op s. Operations op s -> ([BlockItem], [BlockItem])
opsCritical :: ([C.BlockItem], [C.BlockItem])
  }

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

declAllocatedMem :: CompilerM op s [C.BlockItem]
declAllocatedMem :: forall op s. CompilerM op s [BlockItem]
declAllocatedMem = CompilerM op s () -> CompilerM op s [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
collect (CompilerM op s () -> CompilerM op s [BlockItem])
-> CompilerM op s () -> CompilerM op s [BlockItem]
forall a b. (a -> b) -> a -> b
$ ((VName, Space) -> CompilerM op s ())
-> [(VName, Space)] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (VName, Space) -> CompilerM op s ()
forall {op} {s}. (VName, Space) -> CompilerM op s ()
f ([(VName, Space)] -> CompilerM op s ())
-> CompilerM op s [(VName, Space)] -> CompilerM op s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (CompilerState s -> [(VName, Space)])
-> CompilerM op s [(VName, Space)]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> [(VName, Space)]
forall s. CompilerState s -> [(VName, Space)]
compDeclaredMem
  where
    f :: (VName, Space) -> CompilerM op s ()
f (VName
name, Space
space) = do
      Type
ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
memToCType VName
name Space
space
      InitGroup -> CompilerM op s ()
forall op s. InitGroup -> CompilerM op s ()
decl [C.cdecl|$ty:ty $id:name;|]
      VName -> Space -> CompilerM op s ()
forall a op s. ToExp a => a -> Space -> CompilerM op s ()
resetMem VName
name Space
space

data CompilerEnv op s = CompilerEnv
  { forall op s. CompilerEnv op s -> Operations op s
envOperations :: Operations op s,
    -- | Mapping memory blocks to sizes.  These memory blocks are CPU
    -- memory that we know are used in particularly simple ways (no
    -- reference counting necessary).  To cut down on allocator
    -- pressure, we keep these allocations around for a long time, and
    -- record their sizes so we can reuse them if possible (and
    -- realloc() when needed).
    forall op s. CompilerEnv op s -> Map Exp VName
envCachedMem :: M.Map C.Exp VName
  }

contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm], [C.Stm])
contextContents :: forall op s. CompilerM op s ([FieldGroup], [Stm], [Stm])
contextContents = do
  ([Id]
field_names, [Type]
field_types, [Maybe Exp]
field_values, [Maybe (Stm, Stm)]
field_frees) <-
    (CompilerState s
 -> ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)]))
-> CompilerM op s ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)])
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((CompilerState s
  -> ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)]))
 -> CompilerM op s ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)]))
-> (CompilerState s
    -> ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)]))
-> CompilerM op s ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)])
forall a b. (a -> b) -> a -> b
$ [(Id, Type, Maybe Exp, Maybe (Stm, Stm))]
-> ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)])
forall a b c d. [(a, b, c, d)] -> ([a], [b], [c], [d])
unzip4 ([(Id, Type, Maybe Exp, Maybe (Stm, Stm))]
 -> ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)]))
-> (CompilerState s -> [(Id, Type, Maybe Exp, Maybe (Stm, Stm))])
-> CompilerState s
-> ([Id], [Type], [Maybe Exp], [Maybe (Stm, Stm)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DList (Id, Type, Maybe Exp, Maybe (Stm, Stm))
-> [(Id, Type, Maybe Exp, Maybe (Stm, Stm))]
forall a. DList a -> [a]
DL.toList (DList (Id, Type, Maybe Exp, Maybe (Stm, Stm))
 -> [(Id, Type, Maybe Exp, Maybe (Stm, Stm))])
-> (CompilerState s
    -> DList (Id, Type, Maybe Exp, Maybe (Stm, Stm)))
-> CompilerState s
-> [(Id, Type, Maybe Exp, Maybe (Stm, Stm))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerState s -> DList (Id, Type, Maybe Exp, Maybe (Stm, Stm))
forall s.
CompilerState s -> DList (Id, Type, Maybe Exp, Maybe (Stm, Stm))
compCtxFields
  let fields :: [FieldGroup]
fields =
        [ [C.csdecl|$ty:ty $id:name;|]
          | (Id
name, Type
ty) <- [Id] -> [Type] -> [(Id, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
field_names [Type]
field_types
        ]
      init_fields :: [Stm]
init_fields =
        [ [C.cstm|ctx->program->$id:name = $exp:e;|]
          | (Id
name, Just Exp
e) <- [Id] -> [Maybe Exp] -> [(Id, Maybe Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
field_names [Maybe Exp]
field_values
        ]
      ([Stm]
setup, [Stm]
free) = [(Stm, Stm)] -> ([Stm], [Stm])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Stm, Stm)] -> ([Stm], [Stm])) -> [(Stm, Stm)] -> ([Stm], [Stm])
forall a b. (a -> b) -> a -> b
$ [Maybe (Stm, Stm)] -> [(Stm, Stm)]
forall a. [Maybe a] -> [a]
catMaybes [Maybe (Stm, Stm)]
field_frees
  ([FieldGroup], [Stm], [Stm])
-> CompilerM op s ([FieldGroup], [Stm], [Stm])
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([FieldGroup]
fields, [Stm]
init_fields [Stm] -> [Stm] -> [Stm]
forall a. Semigroup a => a -> a -> a
<> [Stm]
setup, [Stm]
free)

generateProgramStruct :: CompilerM op s ()
generateProgramStruct :: forall op s. CompilerM op s ()
generateProgramStruct = do
  ([FieldGroup]
fields, [Stm]
init_fields, [Stm]
free_fields) <- CompilerM op s ([FieldGroup], [Stm], [Stm])
forall op s. CompilerM op s ([FieldGroup], [Stm], [Stm])
contextContents
  (Definition -> CompilerM op s ())
-> [Definition] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_
    Definition -> CompilerM op s ()
forall op s. Definition -> CompilerM op s ()
earlyDecl
    [C.cunit|struct program {
               int dummy;
               $sdecls:fields
             };
             static void setup_program(struct futhark_context* ctx) {
               (void)ctx;
               int error = 0;
               (void)error;
               ctx->program = malloc(sizeof(struct program));
               $stms:init_fields
             }
             static void teardown_program(struct futhark_context *ctx) {
               (void)ctx;
               int error = 0;
               (void)error;
               $stms:free_fields
               free(ctx->program);
             }|]

newtype CompilerM op s a
  = CompilerM (ReaderT (CompilerEnv op s) (State (CompilerState s)) a)
  deriving
    ( (forall a b. (a -> b) -> CompilerM op s a -> CompilerM op s b)
-> (forall a b. a -> CompilerM op s b -> CompilerM op s a)
-> Functor (CompilerM op s)
forall a b. a -> CompilerM op s b -> CompilerM op s a
forall a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a b. a -> CompilerM op s b -> CompilerM op s a
forall op s a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall op s a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
fmap :: forall a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
$c<$ :: forall op s a b. a -> CompilerM op s b -> CompilerM op s a
<$ :: forall a b. a -> CompilerM op s b -> CompilerM op s a
Functor,
      Functor (CompilerM op s)
Functor (CompilerM op s) =>
(forall a. a -> CompilerM op s a)
-> (forall a b.
    CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b)
-> (forall a b c.
    (a -> b -> c)
    -> CompilerM op s a -> CompilerM op s b -> CompilerM op s c)
-> (forall a b.
    CompilerM op s a -> CompilerM op s b -> CompilerM op s b)
-> (forall a b.
    CompilerM op s a -> CompilerM op s b -> CompilerM op s a)
-> Applicative (CompilerM op s)
forall a. a -> CompilerM op s a
forall op s. Functor (CompilerM op s)
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a. a -> CompilerM op s a
forall a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall op s a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s 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 op s a. a -> CompilerM op s a
pure :: forall a. a -> CompilerM op s a
$c<*> :: forall op s a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
<*> :: forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
$cliftA2 :: forall op s a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
liftA2 :: forall a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
$c*> :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
*> :: forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
$c<* :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
<* :: forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
Applicative,
      Applicative (CompilerM op s)
Applicative (CompilerM op s) =>
(forall a b.
 CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b)
-> (forall a b.
    CompilerM op s a -> CompilerM op s b -> CompilerM op s b)
-> (forall a. a -> CompilerM op s a)
-> Monad (CompilerM op s)
forall a. a -> CompilerM op s a
forall op s. Applicative (CompilerM op s)
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
forall op s a. a -> CompilerM op s a
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall op s a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s 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 op s a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
>>= :: forall a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
$c>> :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
>> :: forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
$creturn :: forall op s a. a -> CompilerM op s a
return :: forall a. a -> CompilerM op s a
Monad,
      MonadState (CompilerState s),
      MonadReader (CompilerEnv op s)
    )

instance MonadFreshNames (CompilerM op s) where
  getNameSource :: CompilerM op s VNameSource
getNameSource = (CompilerState s -> VNameSource) -> CompilerM op s VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> VNameSource
forall s. CompilerState s -> VNameSource
compNameSrc
  putNameSource :: VNameSource -> CompilerM op s ()
putNameSource VNameSource
src = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compNameSrc = src}

runCompilerM ::
  Operations op s ->
  VNameSource ->
  s ->
  CompilerM op s a ->
  (a, CompilerState s)
runCompilerM :: forall op s a.
Operations op s
-> VNameSource -> s -> CompilerM op s a -> (a, CompilerState s)
runCompilerM Operations op s
ops VNameSource
src s
userstate (CompilerM ReaderT (CompilerEnv op s) (State (CompilerState s)) a
m) =
  State (CompilerState s) a
-> CompilerState s -> (a, CompilerState s)
forall s a. State s a -> s -> (a, s)
runState
    (ReaderT (CompilerEnv op s) (State (CompilerState s)) a
-> CompilerEnv op s -> State (CompilerState s) a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (CompilerEnv op s) (State (CompilerState s)) a
m (Operations op s -> Map Exp VName -> CompilerEnv op s
forall op s. Operations op s -> Map Exp VName -> CompilerEnv op s
CompilerEnv Operations op s
ops Map Exp VName
forall a. Monoid a => a
mempty))
    (VNameSource -> s -> CompilerState s
forall s. VNameSource -> s -> CompilerState s
newCompilerState VNameSource
src s
userstate)

getUserState :: CompilerM op s s
getUserState :: forall op s. CompilerM op s s
getUserState = (CompilerState s -> s) -> CompilerM op s s
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> s
forall s. CompilerState s -> s
compUserState

modifyUserState :: (s -> s) -> CompilerM op s ()
modifyUserState :: forall s op. (s -> s) -> CompilerM op s ()
modifyUserState s -> s
f = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
compstate ->
  CompilerState s
compstate {compUserState = f $ compUserState compstate}

collect :: CompilerM op s () -> CompilerM op s [C.BlockItem]
collect :: forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
collect CompilerM op s ()
m = ((), [BlockItem]) -> [BlockItem]
forall a b. (a, b) -> b
snd (((), [BlockItem]) -> [BlockItem])
-> CompilerM op s ((), [BlockItem]) -> CompilerM op s [BlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CompilerM op s () -> CompilerM op s ((), [BlockItem])
forall op s a. CompilerM op s a -> CompilerM op s (a, [BlockItem])
collect' CompilerM op s ()
m

collect' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem])
collect' :: forall op s a. CompilerM op s a -> CompilerM op s (a, [BlockItem])
collect' CompilerM op s a
m = do
  DList BlockItem
old <- (CompilerState s -> DList BlockItem)
-> CompilerM op s (DList BlockItem)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> DList BlockItem
forall s. CompilerState s -> DList BlockItem
compItems
  (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compItems = mempty}
  a
x <- CompilerM op s a
m
  DList BlockItem
new <- (CompilerState s -> DList BlockItem)
-> CompilerM op s (DList BlockItem)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> DList BlockItem
forall s. CompilerState s -> DList BlockItem
compItems
  (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compItems = old}
  (a, [BlockItem]) -> CompilerM op s (a, [BlockItem])
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, DList BlockItem -> [BlockItem]
forall a. DList a -> [a]
DL.toList DList BlockItem
new)

-- | Used when we, inside an existing 'CompilerM' action, want to
-- generate code for a new function.  Use this so that the compiler
-- understands that previously declared memory doesn't need to be
-- freed inside this action.
inNewFunction :: CompilerM op s a -> CompilerM op s a
inNewFunction :: forall op s a. CompilerM op s a -> CompilerM op s a
inNewFunction CompilerM op s a
m = do
  [(VName, Space)]
old_mem <- (CompilerState s -> [(VName, Space)])
-> CompilerM op s [(VName, Space)]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> [(VName, Space)]
forall s. CompilerState s -> [(VName, Space)]
compDeclaredMem
  (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compDeclaredMem = mempty}
  a
x <- (CompilerEnv op s -> CompilerEnv op s)
-> CompilerM op s a -> CompilerM op s a
forall a.
(CompilerEnv op s -> CompilerEnv op s)
-> CompilerM op s a -> CompilerM op s a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local CompilerEnv op s -> CompilerEnv op s
forall {op} {s}. CompilerEnv op s -> CompilerEnv op s
noCached CompilerM op s a
m
  (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compDeclaredMem = old_mem}
  a -> CompilerM op s a
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
  where
    noCached :: CompilerEnv op s -> CompilerEnv op s
noCached CompilerEnv op s
env = CompilerEnv op s
env {envCachedMem = mempty}

item :: C.BlockItem -> CompilerM op s ()
item :: forall op s. BlockItem -> CompilerM op s ()
item BlockItem
x = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compItems = DL.snoc (compItems s) x}

items :: [C.BlockItem] -> CompilerM op s ()
items :: forall op s. [BlockItem] -> CompilerM op s ()
items [BlockItem]
xs = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compItems = DL.append (compItems s) (DL.fromList xs)}

fatMemory :: Space -> CompilerM op s Bool
fatMemory :: forall op s. Space -> CompilerM op s Bool
fatMemory ScalarSpace {} = Bool -> CompilerM op s Bool
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
fatMemory Space
_ = (CompilerEnv op s -> Bool) -> CompilerM op s Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((CompilerEnv op s -> Bool) -> CompilerM op s Bool)
-> (CompilerEnv op s -> Bool) -> CompilerM op s Bool
forall a b. (a -> b) -> a -> b
$ Operations op s -> Bool
forall op s. Operations op s -> Bool
opsFatMemory (Operations op s -> Bool)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

cacheMem :: (C.ToExp a) => a -> CompilerM op s (Maybe VName)
cacheMem :: forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
cacheMem a
a = (CompilerEnv op s -> Maybe VName) -> CompilerM op s (Maybe VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((CompilerEnv op s -> Maybe VName) -> CompilerM op s (Maybe VName))
-> (CompilerEnv op s -> Maybe VName)
-> CompilerM op s (Maybe VName)
forall a b. (a -> b) -> a -> b
$ Exp -> Map Exp VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (a -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
C.toExp a
a SrcLoc
forall a. IsLocation a => a
noLoc) (Map Exp VName -> Maybe VName)
-> (CompilerEnv op s -> Map Exp VName)
-> CompilerEnv op s
-> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Map Exp VName
forall op s. CompilerEnv op s -> Map Exp VName
envCachedMem

-- | Construct a publicly visible definition using the specified name
-- as the template.  The first returned definition is put in the
-- header file, and the second is the implementation.  Returns the public
-- name.
publicDef ::
  T.Text ->
  HeaderSection ->
  (T.Text -> (C.Definition, C.Definition)) ->
  CompilerM op s T.Text
publicDef :: forall op s.
Text
-> HeaderSection
-> (Text -> (Definition, Definition))
-> CompilerM op s Text
publicDef Text
s HeaderSection
h Text -> (Definition, Definition)
f = do
  Text
s' <- Text -> CompilerM op s Text
forall op s. Text -> CompilerM op s Text
publicName Text
s
  let (Definition
pub, Definition
priv) = Text -> (Definition, Definition)
f Text
s'
  HeaderSection -> Definition -> CompilerM op s ()
forall op s. HeaderSection -> Definition -> CompilerM op s ()
headerDecl HeaderSection
h Definition
pub
  Definition -> CompilerM op s ()
forall op s. Definition -> CompilerM op s ()
earlyDecl Definition
priv
  Text -> CompilerM op s Text
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
s'

-- | As 'publicDef', but ignores the public name.
publicDef_ ::
  T.Text ->
  HeaderSection ->
  (T.Text -> (C.Definition, C.Definition)) ->
  CompilerM op s ()
publicDef_ :: forall op s.
Text
-> HeaderSection
-> (Text -> (Definition, Definition))
-> CompilerM op s ()
publicDef_ Text
s HeaderSection
h Text -> (Definition, Definition)
f = CompilerM op s Text -> CompilerM op s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (CompilerM op s Text -> CompilerM op s ())
-> CompilerM op s Text -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ Text
-> HeaderSection
-> (Text -> (Definition, Definition))
-> CompilerM op s Text
forall op s.
Text
-> HeaderSection
-> (Text -> (Definition, Definition))
-> CompilerM op s Text
publicDef Text
s HeaderSection
h Text -> (Definition, Definition)
f

headerDecl :: HeaderSection -> C.Definition -> CompilerM op s ()
headerDecl :: forall op s. HeaderSection -> Definition -> CompilerM op s ()
headerDecl HeaderSection
sec Definition
def = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s
    { compHeaderDecls =
        M.unionWith
          (<>)
          (compHeaderDecls s)
          (M.singleton sec (DL.singleton def))
    }

libDecl :: C.Definition -> CompilerM op s ()
libDecl :: forall op s. Definition -> CompilerM op s ()
libDecl Definition
def = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s {compLibDecls = compLibDecls s <> DL.singleton def}

earlyDecl :: C.Definition -> CompilerM op s ()
earlyDecl :: forall op s. Definition -> CompilerM op s ()
earlyDecl Definition
def = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s {compEarlyDecls = compEarlyDecls s <> DL.singleton def}

contextField :: C.Id -> C.Type -> Maybe C.Exp -> CompilerM op s ()
contextField :: forall op s. Id -> Type -> Maybe Exp -> CompilerM op s ()
contextField Id
name Type
ty Maybe Exp
initial = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial, Nothing)}

contextFieldDyn :: C.Id -> C.Type -> C.Stm -> C.Stm -> CompilerM op s ()
contextFieldDyn :: forall op s. Id -> Type -> Stm -> Stm -> CompilerM op s ()
contextFieldDyn Id
name Type
ty Stm
create Stm
free = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, Nothing, Just (create, free))}

onClear :: C.BlockItem -> CompilerM op s ()
onClear :: forall op s. BlockItem -> CompilerM op s ()
onClear BlockItem
x = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s {compClearItems = compClearItems s <> DL.singleton x}

stm :: C.Stm -> CompilerM op s ()
stm :: forall op s. Stm -> CompilerM op s ()
stm Stm
s = BlockItem -> CompilerM op s ()
forall op s. BlockItem -> CompilerM op s ()
item [C.citem|$stm:s|]

stms :: [C.Stm] -> CompilerM op s ()
stms :: forall op s. [Stm] -> CompilerM op s ()
stms = (Stm -> CompilerM op s ()) -> [Stm] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm

decl :: C.InitGroup -> CompilerM op s ()
decl :: forall op s. InitGroup -> CompilerM op s ()
decl InitGroup
x = BlockItem -> CompilerM op s ()
forall op s. BlockItem -> CompilerM op s ()
item [C.citem|$decl:x;|]

-- | Public names must have a consitent prefix.
publicName :: T.Text -> CompilerM op s T.Text
publicName :: forall op s. Text -> CompilerM op s Text
publicName Text
s = Text -> CompilerM op s Text
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> CompilerM op s Text) -> Text -> CompilerM op s Text
forall a b. (a -> b) -> a -> b
$ Text
"futhark_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
s

memToCType :: VName -> Space -> CompilerM op s C.Type
memToCType :: forall op s. VName -> Space -> CompilerM op s Type
memToCType VName
v Space
space = do
  Bool
refcount <- Space -> CompilerM op s Bool
forall op s. Space -> CompilerM op s Bool
fatMemory Space
space
  Bool
cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM op s (Maybe VName) -> CompilerM op s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
cacheMem VName
v
  if Bool
refcount Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
cached
    then Type -> CompilerM op s Type
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> CompilerM op s Type) -> Type -> CompilerM op s Type
forall a b. (a -> b) -> a -> b
$ Space -> Type
fatMemType Space
space
    else Space -> CompilerM op s Type
forall op s. Space -> CompilerM op s Type
rawMemCType Space
space

rawMemCType :: Space -> CompilerM op s C.Type
rawMemCType :: forall op s. Space -> CompilerM op s Type
rawMemCType Space
DefaultSpace = Type -> CompilerM op s Type
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
defaultMemBlockType
rawMemCType (Space [Char]
sid) = CompilerM op s (CompilerM op s Type) -> CompilerM op s Type
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s Type) -> CompilerM op s Type)
-> CompilerM op s (CompilerM op s Type) -> CompilerM op s Type
forall a b. (a -> b) -> a -> b
$ (CompilerEnv op s -> [Char] -> CompilerM op s Type)
-> CompilerM op s ([Char] -> CompilerM op s Type)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Operations op s -> [Char] -> CompilerM op s Type
forall op s. Operations op s -> MemoryType op s
opsMemoryType (Operations op s -> [Char] -> CompilerM op s Type)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> [Char]
-> CompilerM op s Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations) CompilerM op s ([Char] -> CompilerM op s Type)
-> CompilerM op s [Char] -> CompilerM op s (CompilerM op s Type)
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> CompilerM op s [Char]
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
sid
rawMemCType (ScalarSpace [] PrimType
t) =
  Type -> CompilerM op s Type
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cty|$ty:(primTypeToCType t)[1]|]
rawMemCType (ScalarSpace [SubExp]
ds PrimType
t) =
  Type -> CompilerM op s Type
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cty|$ty:(primTypeToCType t)[$exp:(cproduct ds')]|]
  where
    ds' :: [Exp]
ds' = (SubExp -> Exp) -> [SubExp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
`C.toExp` SrcLoc
forall a. IsLocation a => a
noLoc) [SubExp]
ds

fatMemType :: Space -> C.Type
fatMemType :: Space -> Type
fatMemType Space
space =
  [C.cty|struct $id:name|]
  where
    name :: [Char]
name = case Space
space of
      Space [Char]
sid -> [Char]
"memblock_" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
sid
      Space
_ -> [Char]
"memblock"

fatMemSet :: Space -> String
fatMemSet :: Space -> [Char]
fatMemSet (Space [Char]
sid) = [Char]
"memblock_set_" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
sid
fatMemSet Space
_ = [Char]
"memblock_set"

fatMemAlloc :: Space -> String
fatMemAlloc :: Space -> [Char]
fatMemAlloc (Space [Char]
sid) = [Char]
"memblock_alloc_" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
sid
fatMemAlloc Space
_ = [Char]
"memblock_alloc"

fatMemUnRef :: Space -> String
fatMemUnRef :: Space -> [Char]
fatMemUnRef (Space [Char]
sid) = [Char]
"memblock_unref_" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
sid
fatMemUnRef Space
_ = [Char]
"memblock_unref"

rawMem :: VName -> CompilerM op s C.Exp
rawMem :: forall op s. VName -> CompilerM op s Exp
rawMem VName
v = Bool -> VName -> Exp
forall a. ToExp a => Bool -> a -> Exp
rawMem' (Bool -> VName -> Exp)
-> CompilerM op s Bool -> CompilerM op s (VName -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CompilerM op s Bool
forall {op} {s}. CompilerM op s Bool
fat CompilerM op s (VName -> Exp)
-> CompilerM op s VName -> CompilerM op s Exp
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> CompilerM op s VName
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
  where
    fat :: CompilerM op s Bool
fat = (CompilerEnv op s -> Bool -> Bool) -> CompilerM op s (Bool -> Bool)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Bool -> Bool -> Bool
(&&) (Bool -> Bool -> Bool)
-> (CompilerEnv op s -> Bool) -> CompilerEnv op s -> Bool -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Operations op s -> Bool
forall op s. Operations op s -> Bool
opsFatMemory (Operations op s -> Bool)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations) CompilerM op s (Bool -> Bool)
-> CompilerM op s Bool -> CompilerM op s Bool
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Maybe VName -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe VName -> Bool)
-> CompilerM op s (Maybe VName) -> CompilerM op s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
cacheMem VName
v)

rawMem' :: (C.ToExp a) => Bool -> a -> C.Exp
rawMem' :: forall a. ToExp a => Bool -> a -> Exp
rawMem' Bool
True a
e = [C.cexp|$exp:e.mem|]
rawMem' Bool
False a
e = [C.cexp|$exp:e|]

allocRawMem ::
  (C.ToExp a, C.ToExp b, C.ToExp c) =>
  a ->
  b ->
  Space ->
  c ->
  CompilerM op s ()
allocRawMem :: forall a b c op s.
(ToExp a, ToExp b, ToExp c) =>
a -> b -> Space -> c -> CompilerM op s ()
allocRawMem a
dest b
size Space
space c
desc = case Space
space of
  Space [Char]
sid ->
    CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s ()) -> CompilerM op s ())
-> CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$
      (CompilerEnv op s
 -> Exp -> Exp -> Exp -> [Char] -> CompilerM op s ())
-> CompilerM
     op s (Exp -> Exp -> Exp -> [Char] -> CompilerM op s ())
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Operations op s -> Exp -> Exp -> Exp -> [Char] -> CompilerM op s ()
forall op s. Operations op s -> Allocate op s
opsAllocate (Operations op s
 -> Exp -> Exp -> Exp -> [Char] -> CompilerM op s ())
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> Exp
-> Exp
-> Exp
-> [Char]
-> CompilerM op s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations)
        CompilerM op s (Exp -> Exp -> Exp -> [Char] -> CompilerM op s ())
-> CompilerM op s Exp
-> CompilerM op s (Exp -> Exp -> [Char] -> CompilerM op s ())
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s Exp
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cexp|$exp:dest|]
        CompilerM op s (Exp -> Exp -> [Char] -> CompilerM op s ())
-> CompilerM op s Exp
-> CompilerM op s (Exp -> [Char] -> CompilerM op s ())
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s Exp
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cexp|$exp:size|]
        CompilerM op s (Exp -> [Char] -> CompilerM op s ())
-> CompilerM op s Exp
-> CompilerM op s ([Char] -> CompilerM op s ())
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s Exp
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cexp|$exp:desc|]
        CompilerM op s ([Char] -> CompilerM op s ())
-> CompilerM op s [Char] -> CompilerM op s (CompilerM op s ())
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> CompilerM op s [Char]
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
sid
  Space
_ ->
    Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm
      [C.cstm|host_alloc(ctx, (size_t)$exp:size, $exp:desc, (size_t*)&$exp:size, (void*)&$exp:dest);|]

freeRawMem ::
  (C.ToExp a, C.ToExp b, C.ToExp c) =>
  a ->
  b ->
  Space ->
  c ->
  CompilerM op s ()
freeRawMem :: forall a b c op s.
(ToExp a, ToExp b, ToExp c) =>
a -> b -> Space -> c -> CompilerM op s ()
freeRawMem a
mem b
size Space
space c
desc =
  case Space
space of
    Space [Char]
sid -> do
      Deallocate op s
free_mem <- (CompilerEnv op s -> Deallocate op s)
-> CompilerM op s (Deallocate op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Operations op s -> Deallocate op s
forall op s. Operations op s -> Allocate op s
opsDeallocate (Operations op s -> Deallocate op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> Deallocate op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations)
      Deallocate op s
free_mem [C.cexp|$exp:mem|] [C.cexp|$exp:size|] [C.cexp|$exp:desc|] [Char]
sid
    Space
_ ->
      BlockItem -> CompilerM op s ()
forall op s. BlockItem -> CompilerM op s ()
item
        [C.citem|host_free(ctx, (size_t)$exp:size, $exp:desc, (void*)$exp:mem);|]

declMem :: VName -> Space -> CompilerM op s ()
declMem :: forall op s. VName -> Space -> CompilerM op s ()
declMem VName
name Space
space = do
  Bool
cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM op s (Maybe VName) -> CompilerM op s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
cacheMem VName
name
  Bool
fat <- Space -> CompilerM op s Bool
forall op s. Space -> CompilerM op s Bool
fatMemory Space
space
  Bool -> CompilerM op s () -> CompilerM op s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
cached (CompilerM op s () -> CompilerM op s ())
-> CompilerM op s () -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$
    if Bool
fat
      then (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compDeclaredMem = (name, space) : compDeclaredMem s}
      else do
        Type
ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
memToCType VName
name Space
space
        InitGroup -> CompilerM op s ()
forall op s. InitGroup -> CompilerM op s ()
decl [C.cdecl|$ty:ty $id:name;|]

resetMem :: (C.ToExp a) => a -> Space -> CompilerM op s ()
resetMem :: forall a op s. ToExp a => a -> Space -> CompilerM op s ()
resetMem a
mem Space
space = do
  Bool
refcount <- Space -> CompilerM op s Bool
forall op s. Space -> CompilerM op s Bool
fatMemory Space
space
  Bool
cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM op s (Maybe VName) -> CompilerM op s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> CompilerM op s (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
cacheMem a
mem
  if Bool
cached
    then Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$exp:mem = NULL;|]
    else
      Bool -> CompilerM op s () -> CompilerM op s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
refcount (CompilerM op s () -> CompilerM op s ())
-> CompilerM op s () -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$
        Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$exp:mem.references = NULL;|]

setMem :: (C.ToExp a, C.ToExp b) => a -> b -> Space -> CompilerM op s ()
setMem :: forall a b op s.
(ToExp a, ToExp b) =>
a -> b -> Space -> CompilerM op s ()
setMem a
dest b
src Space
space = do
  Bool
refcount <- Space -> CompilerM op s Bool
forall op s. Space -> CompilerM op s Bool
fatMemory Space
space
  let src_s :: [Char]
src_s = Text -> [Char]
T.unpack (Text -> [Char]) -> Text -> [Char]
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
  if Bool
refcount
    then
      Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm
        [C.cstm|if ($id:(fatMemSet space)(ctx, &$exp:dest, &$exp:src,
                                               $string:src_s) != 0) {
                       return 1;
                     }|]
    else case Space
space of
      ScalarSpace [SubExp]
ds PrimType
_ -> do
        VName
i' <- [Char] -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"i"
        let i :: SrcLoc -> Id
i = VName -> SrcLoc -> Id
forall a. ToIdent a => a -> SrcLoc -> Id
C.toIdent VName
i'
            it :: Type
it = PrimType -> Type
primTypeToCType (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32
            ds' :: [Exp]
ds' = (SubExp -> Exp) -> [SubExp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
`C.toExp` SrcLoc
forall a. IsLocation a => a
noLoc) [SubExp]
ds
            bound :: Exp
bound = [Exp] -> Exp
cproduct [Exp]
ds'
        Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm
          [C.cstm|for ($ty:it $id:i = 0; $id:i < $exp:bound; $id:i++) {
                            $exp:dest[$id:i] = $exp:src[$id:i];
                  }|]
      Space
_ -> Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$exp:dest = $exp:src;|]

unRefMem :: (C.ToExp a) => a -> Space -> CompilerM op s ()
unRefMem :: forall a op s. ToExp a => a -> Space -> CompilerM op s ()
unRefMem a
mem Space
space = do
  Bool
refcount <- Space -> CompilerM op s Bool
forall op s. Space -> CompilerM op s Bool
fatMemory Space
space
  Bool
cached <- Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool)
-> CompilerM op s (Maybe VName) -> CompilerM op s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> CompilerM op s (Maybe VName)
forall a op s. ToExp a => a -> CompilerM op s (Maybe VName)
cacheMem a
mem
  let mem_s :: [Char]
mem_s = Text -> [Char]
T.unpack (Text -> [Char]) -> Text -> [Char]
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
  Bool -> CompilerM op s () -> CompilerM op s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool
refcount Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
cached) (CompilerM op s () -> CompilerM op s ())
-> CompilerM op s () -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$
    Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm
      [C.cstm|if ($id:(fatMemUnRef space)(ctx, &$exp:mem, $string:mem_s) != 0) {
                  return 1;
                }|]

allocMem ::
  (C.ToExp a, C.ToExp b) =>
  a ->
  b ->
  Space ->
  C.Stm ->
  CompilerM op s ()
allocMem :: forall a b op s.
(ToExp a, ToExp b) =>
a -> b -> Space -> Stm -> CompilerM op s ()
allocMem a
mem b
size Space
space Stm
on_failure = do
  Bool
refcount <- Space -> CompilerM op s Bool
forall op s. Space -> CompilerM op s Bool
fatMemory Space
space
  let mem_s :: [Char]
mem_s = Text -> [Char]
T.unpack (Text -> [Char]) -> Text -> [Char]
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
  if Bool
refcount
    then
      Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm
        [C.cstm|if ($id:(fatMemAlloc space)(ctx, &$exp:mem, $exp:size,
                                                 $string:mem_s)) {
                       $stm:on_failure
                     }|]
    else do
      a -> b -> Space -> [Char] -> CompilerM op s ()
forall a b c op s.
(ToExp a, ToExp b, ToExp c) =>
a -> b -> Space -> c -> CompilerM op s ()
freeRawMem a
mem b
size Space
space [Char]
mem_s
      a -> b -> Space -> Exp -> CompilerM op s ()
forall a b c op s.
(ToExp a, ToExp b, ToExp c) =>
a -> b -> Space -> c -> CompilerM op s ()
allocRawMem a
mem b
size Space
space [C.cexp|desc|]

copyMemoryDefaultSpace ::
  C.Exp ->
  C.Exp ->
  C.Exp ->
  C.Exp ->
  C.Exp ->
  CompilerM op s ()
copyMemoryDefaultSpace :: forall op s. Exp -> Exp -> Exp -> Exp -> Exp -> CompilerM op s ()
copyMemoryDefaultSpace Exp
destmem Exp
destidx Exp
srcmem Exp
srcidx Exp
nbytes =
  Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm
    [C.cstm|if ($exp:nbytes > 0) {
              memmove($exp:destmem + $exp:destidx,
                      $exp:srcmem + $exp:srcidx,
                      $exp:nbytes);
            }|]

cachingMemory ::
  M.Map VName Space ->
  ([C.BlockItem] -> [C.Stm] -> CompilerM op s a) ->
  CompilerM op s a
cachingMemory :: forall op s a.
Map VName Space
-> ([BlockItem] -> [Stm] -> CompilerM op s a) -> CompilerM op s a
cachingMemory Map VName Space
lexical [BlockItem] -> [Stm] -> CompilerM op s a
f = do
  -- We only consider lexical 'DefaultSpace' memory blocks to be
  -- cached.  This is not a deep technical restriction, but merely a
  -- heuristic based on GPU memory usually involving larger
  -- allocations, that do not suffer from the overhead of reference
  -- counting.  Beware: there is code elsewhere in codegen that
  -- assumes lexical memory is DefaultSpace too.
  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

  [(VName, VName)]
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
    VName
size <- [Char] -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> CompilerM op s VName) -> [Char] -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
mem [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Char]
"_cached_size"
    (VName, VName) -> CompilerM op s (VName, VName)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
mem, VName
size)

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

      declCached :: (a, a) -> [BlockItem]
declCached (a
mem, a
size) =
        [ [C.citem|typename int64_t $id:size = 0;|],
          [C.citem|$ty:defaultMemBlockType $id:mem = NULL;|]
        ]

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

  (CompilerEnv op s -> CompilerEnv op s)
-> CompilerM op s a -> CompilerM op s a
forall a.
(CompilerEnv op s -> CompilerEnv op s)
-> CompilerM op s a -> CompilerM op s a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local CompilerEnv op s -> CompilerEnv op s
forall {op} {s}. CompilerEnv op s -> CompilerEnv op s
lexMem (CompilerM op s a -> CompilerM op s a)
-> CompilerM op s a -> CompilerM op s a
forall a b. (a -> b) -> a -> b
$ [BlockItem] -> [Stm] -> CompilerM op s a
f (((VName, VName) -> [BlockItem]) -> [(VName, VName)] -> [BlockItem]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (VName, VName) -> [BlockItem]
forall {a} {a}. (ToIdent a, ToIdent a) => (a, a) -> [BlockItem]
declCached [(VName, VName)]
cached') (((VName, VName) -> Stm) -> [(VName, VName)] -> [Stm]
forall a b. (a -> b) -> [a] -> [b]
map (VName, VName) -> Stm
forall {a} {b}. ToIdent a => (a, b) -> Stm
freeCached [(VName, VName)]
cached')

derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp
derefPointer :: Exp -> Exp -> Type -> Exp
derefPointer Exp
ptr Exp
i Type
res_t =
  [C.cexp|(($ty:res_t)$exp:ptr)[$exp:i]|]

volQuals :: Volatility -> [C.TypeQual]
volQuals :: Volatility -> [TypeQual]
volQuals Volatility
Volatile = [C.ctyquals|volatile|]
volQuals Volatility
Nonvolatile = []

writeScalarPointerWithQuals :: PointerQuals -> WriteScalar op s
writeScalarPointerWithQuals :: forall op s. PointerQuals -> WriteScalar op s
writeScalarPointerWithQuals PointerQuals
quals_f Exp
dest Exp
i Type
elemtype [Char]
space Volatility
vol Exp
v = do
  let quals' :: [TypeQual]
quals' = Volatility -> [TypeQual]
volQuals Volatility
vol [TypeQual] -> [TypeQual] -> [TypeQual]
forall a. [a] -> [a] -> [a]
++ PointerQuals
quals_f [Char]
space
      deref :: Exp
deref = Exp -> Exp -> Type -> Exp
derefPointer Exp
dest Exp
i [C.cty|$tyquals:quals' $ty:elemtype*|]
  Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$exp:deref = $exp:v;|]

readScalarPointerWithQuals :: PointerQuals -> ReadScalar op s
readScalarPointerWithQuals :: forall op s. PointerQuals -> ReadScalar op s
readScalarPointerWithQuals PointerQuals
quals_f Exp
dest Exp
i Type
elemtype [Char]
space Volatility
vol = do
  let quals' :: [TypeQual]
quals' = Volatility -> [TypeQual]
volQuals Volatility
vol [TypeQual] -> [TypeQual] -> [TypeQual]
forall a. [a] -> [a] -> [a]
++ PointerQuals
quals_f [Char]
space
  Exp -> CompilerM op s Exp
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> CompilerM op s Exp) -> Exp -> CompilerM op s Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Type -> Exp
derefPointer Exp
dest Exp
i [C.cty|$tyquals:quals' $ty:elemtype*|]

criticalSection :: Operations op s -> [C.BlockItem] -> [C.BlockItem]
criticalSection :: forall op s. Operations op s -> [BlockItem] -> [BlockItem]
criticalSection Operations op s
ops [BlockItem]
x =
  [C.citems|lock_lock(&ctx->lock);
            $items:(fst (opsCritical ops))
            $items:x
            $items:(snd (opsCritical ops))
            lock_unlock(&ctx->lock);
           |]

-- | The generated code must define a context struct with this name.
contextType :: CompilerM op s C.Type
contextType :: forall op s. CompilerM op s Type
contextType = do
  Text
name <- Text -> CompilerM op s Text
forall op s. Text -> CompilerM op s Text
publicName Text
"context"
  Type -> CompilerM op s Type
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cty|struct $id:name|]

-- | The generated code must define a configuration struct with this
-- name.
configType :: CompilerM op s C.Type
configType :: forall op s. CompilerM op s Type
configType = do
  Text
name <- Text -> CompilerM op s Text
forall op s. Text -> CompilerM op s Text
publicName Text
"context_config"
  Type -> CompilerM op s Type
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cty|struct $id:name|]