{-# LANGUAGE QuasiQuotes #-}

-- | Generate the entry point packing/unpacking code.
module Futhark.CodeGen.Backends.GenericC.EntryPoints
  ( onEntryPoint,
  )
where

import Control.Monad
import Control.Monad.Reader (asks)
import Data.Maybe
import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericC.Monad
import Futhark.CodeGen.Backends.GenericC.Types (opaqueToCType, valueTypeToCType)
import Futhark.CodeGen.ImpCode
import Futhark.Manifest qualified as Manifest
import Language.C.Quote.OpenCL qualified as C
import Language.C.Syntax qualified as C

valueDescToType :: ValueDesc -> ValueType
valueDescToType :: ValueDesc -> ValueType
valueDescToType (ScalarValue PrimType
pt Signedness
signed VName
_) =
  Signedness -> Rank -> PrimType -> ValueType
ValueType Signedness
signed (Int -> Rank
Rank Int
0) PrimType
pt
valueDescToType (ArrayValue VName
_ Space
_ PrimType
pt Signedness
signed [DimSize]
shape) =
  Signedness -> Rank -> PrimType -> ValueType
ValueType Signedness
signed (Int -> Rank
Rank ([DimSize] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
shape)) PrimType
pt

allTrue :: [C.Exp] -> C.Exp
allTrue :: [Exp] -> Exp
allTrue [] = [C.cexp|true|]
allTrue [Exp
x] = Exp
x
allTrue (Exp
x : [Exp]
xs) = [C.cexp|$exp:x && $exp:(allTrue xs)|]

prepareEntryInputs ::
  [ExternalValue] ->
  CompilerM op s ([(C.Param, Maybe C.Exp)], [C.BlockItem])
prepareEntryInputs :: forall op s.
[ExternalValue]
-> CompilerM op s ([(Param, Maybe Exp)], [BlockItem])
prepareEntryInputs [ExternalValue]
args = CompilerM op s [(Param, Maybe Exp)]
-> CompilerM op s ([(Param, Maybe Exp)], [BlockItem])
forall op s a. CompilerM op s a -> CompilerM op s (a, [BlockItem])
collect' (CompilerM op s [(Param, Maybe Exp)]
 -> CompilerM op s ([(Param, Maybe Exp)], [BlockItem]))
-> CompilerM op s [(Param, Maybe Exp)]
-> CompilerM op s ([(Param, Maybe Exp)], [BlockItem])
forall a b. (a -> b) -> a -> b
$ (Int -> ExternalValue -> CompilerM op s (Param, Maybe Exp))
-> [Int] -> [ExternalValue] -> CompilerM op s [(Param, Maybe Exp)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Int -> ExternalValue -> CompilerM op s (Param, Maybe Exp)
forall {p} {op} {s}.
Show p =>
p -> ExternalValue -> CompilerM op s (Param, Maybe Exp)
prepare [(Int
0 :: Int) ..] [ExternalValue]
args
  where
    arg_names :: Names
arg_names = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (ExternalValue -> [VName]) -> [ExternalValue] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ExternalValue -> [VName]
evNames [ExternalValue]
args
    evNames :: ExternalValue -> [VName]
evNames (OpaqueValue Name
_ [ValueDesc]
vds) = (ValueDesc -> VName) -> [ValueDesc] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ValueDesc -> VName
vdName [ValueDesc]
vds
    evNames (TransparentValue ValueDesc
vd) = [ValueDesc -> VName
vdName ValueDesc
vd]
    vdName :: ValueDesc -> VName
vdName (ArrayValue VName
v Space
_ PrimType
_ Signedness
_ [DimSize]
_) = VName
v
    vdName (ScalarValue PrimType
_ Signedness
_ VName
v) = VName
v

    prepare :: p -> ExternalValue -> CompilerM op s (Param, Maybe Exp)
prepare p
pno (TransparentValue ValueDesc
vd) = do
      let pname :: FilePath
pname = FilePath
"in" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ p -> FilePath
forall a. Show a => a -> FilePath
show p
pno
      (Type
ty, [Exp]
check) <- Publicness -> Exp -> ValueDesc -> CompilerM op s (Type, [Exp])
forall {a} {op} {s}.
ToExp a =>
Publicness -> a -> ValueDesc -> CompilerM op s (Type, [Exp])
prepareValue Publicness
Public [C.cexp|$id:pname|] ValueDesc
vd
      (Param, Maybe Exp) -> CompilerM op s (Param, Maybe Exp)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( [C.cparam|const $ty:ty $id:pname|],
          if [Exp] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Exp]
check then Maybe Exp
forall a. Maybe a
Nothing else Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
allTrue [Exp]
check
        )
    prepare p
pno (OpaqueValue Name
desc [ValueDesc]
vds) = do
      Type
ty <- Name -> CompilerM op s Type
forall op s. Name -> CompilerM op s Type
opaqueToCType Name
desc
      let pname :: FilePath
pname = FilePath
"in" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ p -> FilePath
forall a. Show a => a -> FilePath
show p
pno
          field :: Int -> ValueDesc -> Exp
field Int
i ScalarValue {} = [C.cexp|$id:pname->$id:(tupleField i)|]
          field Int
i ArrayValue {} = [C.cexp|$id:pname->$id:(tupleField i)|]
      [[Exp]]
checks <- ((Type, [Exp]) -> [Exp]) -> [(Type, [Exp])] -> [[Exp]]
forall a b. (a -> b) -> [a] -> [b]
map (Type, [Exp]) -> [Exp]
forall a b. (a, b) -> b
snd ([(Type, [Exp])] -> [[Exp]])
-> CompilerM op s [(Type, [Exp])] -> CompilerM op s [[Exp]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> ValueDesc -> CompilerM op s (Type, [Exp]))
-> [Exp] -> [ValueDesc] -> CompilerM op s [(Type, [Exp])]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (Publicness -> Exp -> ValueDesc -> CompilerM op s (Type, [Exp])
forall {a} {op} {s}.
ToExp a =>
Publicness -> a -> ValueDesc -> CompilerM op s (Type, [Exp])
prepareValue Publicness
Private) ((Int -> ValueDesc -> Exp) -> [Int] -> [ValueDesc] -> [Exp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> ValueDesc -> Exp
field [Int
0 ..] [ValueDesc]
vds) [ValueDesc]
vds
      (Param, Maybe Exp) -> CompilerM op s (Param, Maybe Exp)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( [C.cparam|const $ty:ty *$id:pname|],
          if ([Exp] -> Bool) -> [[Exp]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all [Exp] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Exp]]
checks
            then Maybe Exp
forall a. Maybe a
Nothing
            else Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
allTrue ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ [[Exp]] -> [Exp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Exp]]
checks
        )

    prepareValue :: Publicness -> a -> ValueDesc -> CompilerM op s (Type, [Exp])
prepareValue Publicness
_ a
src (ScalarValue PrimType
pt Signedness
signed VName
name) = do
      let pt' :: Type
pt' = Signedness -> PrimType -> Type
primAPIType Signedness
signed PrimType
pt
          src' :: Exp
src' = PrimType -> Exp -> Exp
fromStorage PrimType
pt (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ a -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
C.toExp a
src SrcLoc
forall a. Monoid a => a
mempty
      Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$id:name = $exp:src';|]
      (Type, [Exp]) -> CompilerM op s (Type, [Exp])
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
pt', [])
    prepareValue Publicness
pub a
src vd :: ValueDesc
vd@(ArrayValue VName
mem Space
_ PrimType
_ Signedness
_ [DimSize]
shape) = do
      Type
ty <- Publicness -> ValueType -> CompilerM op s Type
forall op s. Publicness -> ValueType -> CompilerM op s Type
valueTypeToCType Publicness
pub (ValueType -> CompilerM op s Type)
-> ValueType -> CompilerM op s Type
forall a b. (a -> b) -> a -> b
$ ValueDesc -> ValueType
valueDescToType ValueDesc
vd

      Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$exp:mem = $exp:src->mem;|]

      let rank :: Int
rank = [DimSize] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
shape
          maybeCopyDim :: DimSize -> a -> (Maybe Stm, Exp)
maybeCopyDim (Var VName
d) a
i
            | VName
d VName -> Names -> Bool
`notNameIn` Names
arg_names =
                ( Stm -> Maybe Stm
forall a. a -> Maybe a
Just [C.cstm|$id:d = $exp:src->shape[$int:i];|],
                  [C.cexp|$id:d == $exp:src->shape[$int:i]|]
                )
          maybeCopyDim DimSize
x a
i =
            ( Maybe Stm
forall a. Maybe a
Nothing,
              [C.cexp|$exp:x == $exp:src->shape[$int:i]|]
            )

      let ([Maybe Stm]
sets, [Exp]
checks) =
            [(Maybe Stm, Exp)] -> ([Maybe Stm], [Exp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Maybe Stm, Exp)] -> ([Maybe Stm], [Exp]))
-> [(Maybe Stm, Exp)] -> ([Maybe Stm], [Exp])
forall a b. (a -> b) -> a -> b
$ (DimSize -> Int -> (Maybe Stm, Exp))
-> [DimSize] -> [Int] -> [(Maybe Stm, Exp)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith DimSize -> Int -> (Maybe Stm, Exp)
forall {a}.
(Show a, Integral a) =>
DimSize -> a -> (Maybe Stm, Exp)
maybeCopyDim [DimSize]
shape [Int
0 .. Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      [Stm] -> CompilerM op s ()
forall op s. [Stm] -> CompilerM op s ()
stms ([Stm] -> CompilerM op s ()) -> [Stm] -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ [Maybe Stm] -> [Stm]
forall a. [Maybe a] -> [a]
catMaybes [Maybe Stm]
sets

      (Type, [Exp]) -> CompilerM op s (Type, [Exp])
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([C.cty|$ty:ty*|], [Exp]
checks)

prepareEntryOutputs :: [ExternalValue] -> CompilerM op s ([C.Param], [C.BlockItem])
prepareEntryOutputs :: forall op s.
[ExternalValue] -> CompilerM op s ([Param], [BlockItem])
prepareEntryOutputs = CompilerM op s [Param] -> CompilerM op s ([Param], [BlockItem])
forall op s a. CompilerM op s a -> CompilerM op s (a, [BlockItem])
collect' (CompilerM op s [Param] -> CompilerM op s ([Param], [BlockItem]))
-> ([ExternalValue] -> CompilerM op s [Param])
-> [ExternalValue]
-> CompilerM op s ([Param], [BlockItem])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> ExternalValue -> CompilerM op s Param)
-> [Int] -> [ExternalValue] -> CompilerM op s [Param]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Int -> ExternalValue -> CompilerM op s Param
forall {p} {op} {s}.
Show p =>
p -> ExternalValue -> CompilerM op s Param
prepare [(Int
0 :: Int) ..]
  where
    prepare :: p -> ExternalValue -> CompilerM op s Param
prepare p
pno (TransparentValue ValueDesc
vd) = do
      let pname :: FilePath
pname = FilePath
"out" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ p -> FilePath
forall a. Show a => a -> FilePath
show p
pno
      Type
ty <- Publicness -> ValueType -> CompilerM op s Type
forall op s. Publicness -> ValueType -> CompilerM op s Type
valueTypeToCType Publicness
Public (ValueType -> CompilerM op s Type)
-> ValueType -> CompilerM op s Type
forall a b. (a -> b) -> a -> b
$ ValueDesc -> ValueType
valueDescToType ValueDesc
vd

      case ValueDesc
vd of
        ArrayValue {} -> do
          Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|]
          Exp -> ValueDesc -> CompilerM op s ()
forall {a} {op} {s}. ToExp a => a -> ValueDesc -> CompilerM op s ()
prepareValue [C.cexp|*$id:pname|] ValueDesc
vd
          Param -> CompilerM op s Param
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cparam|$ty:ty **$id:pname|]
        ScalarValue {} -> do
          Exp -> ValueDesc -> CompilerM op s ()
forall {a} {op} {s}. ToExp a => a -> ValueDesc -> CompilerM op s ()
prepareValue [C.cexp|*$id:pname|] ValueDesc
vd
          Param -> CompilerM op s Param
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cparam|$ty:ty *$id:pname|]
    prepare p
pno (OpaqueValue Name
desc [ValueDesc]
vds) = do
      let pname :: FilePath
pname = FilePath
"out" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ p -> FilePath
forall a. Show a => a -> FilePath
show p
pno
      Type
ty <- Name -> CompilerM op s Type
forall op s. Name -> CompilerM op s Type
opaqueToCType Name
desc
      [Type]
vd_ts <- (ValueDesc -> CompilerM op s Type)
-> [ValueDesc] -> CompilerM op s [Type]
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 (Publicness -> ValueType -> CompilerM op s Type
forall op s. Publicness -> ValueType -> CompilerM op s Type
valueTypeToCType Publicness
Private (ValueType -> CompilerM op s Type)
-> (ValueDesc -> ValueType) -> ValueDesc -> CompilerM op s Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueDesc -> ValueType
valueDescToType) [ValueDesc]
vds

      Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|]

      [(Int, Type, ValueDesc)]
-> ((Int, Type, ValueDesc) -> CompilerM op s ())
-> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [Type] -> [ValueDesc] -> [(Int, Type, ValueDesc)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Int
0 ..] [Type]
vd_ts [ValueDesc]
vds) (((Int, Type, ValueDesc) -> CompilerM op s ())
 -> CompilerM op s ())
-> ((Int, Type, ValueDesc) -> CompilerM op s ())
-> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, Type
ct, ValueDesc
vd) -> do
        let field :: Exp
field = [C.cexp|((*$id:pname)->$id:(tupleField i))|]
        case ValueDesc
vd of
          ScalarValue {} -> () -> CompilerM op s ()
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          ArrayValue {} -> do
            Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|assert(($exp:field = ($ty:ct*) malloc(sizeof($ty:ct))) != NULL);|]
        Exp -> ValueDesc -> CompilerM op s ()
forall {a} {op} {s}. ToExp a => a -> ValueDesc -> CompilerM op s ()
prepareValue Exp
field ValueDesc
vd

      Param -> CompilerM op s Param
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cparam|$ty:ty **$id:pname|]

    prepareValue :: a -> ValueDesc -> CompilerM op s ()
prepareValue a
dest (ScalarValue PrimType
t Signedness
_ VName
name) =
      let name' :: Exp
name' = PrimType -> Exp -> Exp
toStorage PrimType
t (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ VName -> SrcLoc -> Exp
forall a. ToExp a => a -> SrcLoc -> Exp
C.toExp VName
name SrcLoc
forall a. Monoid a => a
mempty
       in Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$exp:dest = $exp:name';|]
    prepareValue a
dest (ArrayValue VName
mem Space
_ PrimType
_ Signedness
_ [DimSize]
shape) = do
      Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|$exp:dest->mem = $id:mem;|]

      let rank :: Int
rank = [DimSize] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
shape
          maybeCopyDim :: DimSize -> a -> Stm
maybeCopyDim (Constant PrimValue
x) a
i =
            [C.cstm|$exp:dest->shape[$int:i] = $exp:x;|]
          maybeCopyDim (Var VName
d) a
i =
            [C.cstm|$exp:dest->shape[$int:i] = $id:d;|]
      [Stm] -> CompilerM op s ()
forall op s. [Stm] -> CompilerM op s ()
stms ([Stm] -> CompilerM op s ()) -> [Stm] -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ (DimSize -> Int -> Stm) -> [DimSize] -> [Int] -> [Stm]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith DimSize -> Int -> Stm
forall {a}. (Show a, Integral a) => DimSize -> a -> Stm
maybeCopyDim [DimSize]
shape [Int
0 .. Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

entryName :: Name -> T.Text
entryName :: Name -> Text
entryName = (Text
"entry_" <>) (Text -> Text) -> (Name -> Text) -> Name -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
escapeName (Text -> Text) -> (Name -> Text) -> Name -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Text
nameToText

onEntryPoint ::
  [C.BlockItem] ->
  [Name] ->
  Name ->
  Function op ->
  CompilerM op s (Maybe (C.Definition, (T.Text, Manifest.EntryPoint)))
onEntryPoint :: forall op s.
[BlockItem]
-> [Name]
-> Name
-> Function op
-> CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
onEntryPoint [BlockItem]
_ [Name]
_ Name
_ (Function Maybe EntryPoint
Nothing [Param]
_ [Param]
_ Code op
_) = Maybe (Definition, (Text, EntryPoint))
-> CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Definition, (Text, EntryPoint))
forall a. Maybe a
Nothing
onEntryPoint [BlockItem]
get_consts [Name]
relevant_params Name
fname (Function (Just (EntryPoint Name
ename [(Uniqueness, ExternalValue)]
results [((Name, Uniqueness), ExternalValue)]
args)) [Param]
outputs [Param]
inputs Code op
_) = CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
-> CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
forall op s a. CompilerM op s a -> CompilerM op s a
inNewFunction (CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
 -> CompilerM op s (Maybe (Definition, (Text, EntryPoint))))
-> CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
-> CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
forall a b. (a -> b) -> a -> b
$ do
  let out_args :: [Exp]
out_args = (Param -> Exp) -> [Param] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (\Param
p -> [C.cexp|&$id:(paramName p)|]) [Param]
outputs
      in_args :: [Exp]
in_args = (Param -> Exp) -> [Param] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (\Param
p -> [C.cexp|$id:(paramName p)|]) [Param]
inputs

  [BlockItem]
inputdecls <- 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
$ (Param -> CompilerM op s ()) -> [Param] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Param -> CompilerM op s ()
forall {op} {s}. Param -> CompilerM op s ()
stubParam [Param]
inputs
  [BlockItem]
outputdecls <- 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
$ (Param -> CompilerM op s ()) -> [Param] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Param -> CompilerM op s ()
forall {op} {s}. Param -> CompilerM op s ()
stubParam [Param]
outputs
  [BlockItem]
decl_mem <- CompilerM op s [BlockItem]
forall op s. CompilerM op s [BlockItem]
declAllocatedMem

  Text
entry_point_function_name <- Text -> CompilerM op s Text
forall op s. Text -> CompilerM op s Text
publicName (Text -> CompilerM op s Text) -> Text -> CompilerM op s Text
forall a b. (a -> b) -> a -> b
$ Name -> Text
entryName Name
ename

  ([(Param, Maybe Exp)]
inputs', [BlockItem]
unpack_entry_inputs) <- [ExternalValue]
-> CompilerM op s ([(Param, Maybe Exp)], [BlockItem])
forall op s.
[ExternalValue]
-> CompilerM op s ([(Param, Maybe Exp)], [BlockItem])
prepareEntryInputs ([ExternalValue]
 -> CompilerM op s ([(Param, Maybe Exp)], [BlockItem]))
-> [ExternalValue]
-> CompilerM op s ([(Param, Maybe Exp)], [BlockItem])
forall a b. (a -> b) -> a -> b
$ (((Name, Uniqueness), ExternalValue) -> ExternalValue)
-> [((Name, Uniqueness), ExternalValue)] -> [ExternalValue]
forall a b. (a -> b) -> [a] -> [b]
map ((Name, Uniqueness), ExternalValue) -> ExternalValue
forall a b. (a, b) -> b
snd [((Name, Uniqueness), ExternalValue)]
args
  let ([Param]
entry_point_input_params, [Maybe Exp]
entry_point_input_checks) = [(Param, Maybe Exp)] -> ([Param], [Maybe Exp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param, Maybe Exp)]
inputs'

  ([Param]
entry_point_output_params, [BlockItem]
pack_entry_outputs) <-
    [ExternalValue] -> CompilerM op s ([Param], [BlockItem])
forall op s.
[ExternalValue] -> CompilerM op s ([Param], [BlockItem])
prepareEntryOutputs ([ExternalValue] -> CompilerM op s ([Param], [BlockItem]))
-> [ExternalValue] -> CompilerM op s ([Param], [BlockItem])
forall a b. (a -> b) -> a -> b
$ ((Uniqueness, ExternalValue) -> ExternalValue)
-> [(Uniqueness, ExternalValue)] -> [ExternalValue]
forall a b. (a -> b) -> [a] -> [b]
map (Uniqueness, ExternalValue) -> ExternalValue
forall a b. (a, b) -> b
snd [(Uniqueness, ExternalValue)]
results

  Type
ctx_ty <- CompilerM op s Type
forall op s. CompilerM op s Type
contextType

  HeaderSection -> Definition -> CompilerM op s ()
forall op s. HeaderSection -> Definition -> CompilerM op s ()
headerDecl
    HeaderSection
EntryDecl
    [C.cedecl|int $id:entry_point_function_name
                                     ($ty:ctx_ty *ctx,
                                      $params:entry_point_output_params,
                                      $params:entry_point_input_params);|]

  let checks :: [Exp]
checks = [Maybe Exp] -> [Exp]
forall a. [Maybe a] -> [a]
catMaybes [Maybe Exp]
entry_point_input_checks
      check_input :: [BlockItem]
check_input =
        if [Exp] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Exp]
checks
          then []
          else
            [C.citems|
         if (!($exp:(allTrue (catMaybes entry_point_input_checks)))) {
           ret = 1;
           set_error(ctx, msgprintf("Error: entry point arguments have invalid sizes.\n"));
         }|]

      critical :: [BlockItem]
critical =
        [C.citems|
         $items:decl_mem
         $items:unpack_entry_inputs
         $items:check_input
         if (ret == 0) {
           ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args);
           if (ret == 0) {
             $items:get_consts

             $items:pack_entry_outputs
           }
         }
        |]

  Operations op s
ops <- (CompilerEnv op s -> Operations op s)
-> CompilerM op s (Operations op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

  let cdef :: Definition
cdef =
        [C.cedecl|
       int $id:entry_point_function_name
           ($ty:ctx_ty *ctx,
            $params:entry_point_output_params,
            $params:entry_point_input_params) {
         $items:inputdecls
         $items:outputdecls

         int ret = 0;

         $items:(criticalSection ops critical)

         return ret;
       }
       |]

      manifest :: EntryPoint
manifest =
        Manifest.EntryPoint
          { entryPointCFun :: Text
Manifest.entryPointCFun = Text
entry_point_function_name,
            entryPointTuningParams :: [Text]
Manifest.entryPointTuningParams = (Name -> Text) -> [Name] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Text
nameToText [Name]
relevant_params,
            -- Note that our convention about what is "input/output"
            -- and what is "results/args" is different between the
            -- manifest and ImpCode.
            entryPointOutputs :: [Output]
Manifest.entryPointOutputs = ((Uniqueness, ExternalValue) -> Output)
-> [(Uniqueness, ExternalValue)] -> [Output]
forall a b. (a -> b) -> [a] -> [b]
map (Uniqueness, ExternalValue) -> Output
outputManifest [(Uniqueness, ExternalValue)]
results,
            entryPointInputs :: [Input]
Manifest.entryPointInputs = (((Name, Uniqueness), ExternalValue) -> Input)
-> [((Name, Uniqueness), ExternalValue)] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map ((Name, Uniqueness), ExternalValue) -> Input
inputManifest [((Name, Uniqueness), ExternalValue)]
args
          }

  Maybe (Definition, (Text, EntryPoint))
-> CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Definition, (Text, EntryPoint))
 -> CompilerM op s (Maybe (Definition, (Text, EntryPoint))))
-> Maybe (Definition, (Text, EntryPoint))
-> CompilerM op s (Maybe (Definition, (Text, EntryPoint)))
forall a b. (a -> b) -> a -> b
$ (Definition, (Text, EntryPoint))
-> Maybe (Definition, (Text, EntryPoint))
forall a. a -> Maybe a
Just (Definition
cdef, (Name -> Text
nameToText Name
ename, EntryPoint
manifest))
  where
    stubParam :: Param -> CompilerM op s ()
stubParam (MemParam VName
name Space
space) =
      VName -> Space -> CompilerM op s ()
forall op s. VName -> Space -> CompilerM op s ()
declMem VName
name Space
space
    stubParam (ScalarParam VName
name PrimType
ty) = do
      let ty' :: Type
ty' = PrimType -> Type
primTypeToCType PrimType
ty
      InitGroup -> CompilerM op s ()
forall op s. InitGroup -> CompilerM op s ()
decl [C.cdecl|$ty:ty' $id:name = $exp:(blankPrimValue ty);|]

    vdType :: ExternalValue -> Text
vdType (TransparentValue (ScalarValue PrimType
pt Signedness
signed VName
_)) =
      Bool -> PrimType -> Text
prettySigned (Signedness
signed Signedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
== Signedness
Unsigned) PrimType
pt
    vdType (TransparentValue (ArrayValue VName
_ Space
_ PrimType
pt Signedness
signed [DimSize]
shape)) =
      [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat (Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate ([DimSize] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
shape) Text
"[]")
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Bool -> PrimType -> Text
prettySigned (Signedness
signed Signedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
== Signedness
Unsigned) PrimType
pt
    vdType (OpaqueValue Name
name [ValueDesc]
_) =
      Name -> Text
nameToText Name
name

    outputManifest :: (Uniqueness, ExternalValue) -> Output
outputManifest (Uniqueness
u, ExternalValue
vd) =
      Manifest.Output
        { outputType :: Text
Manifest.outputType = ExternalValue -> Text
vdType ExternalValue
vd,
          outputUnique :: Bool
Manifest.outputUnique = Uniqueness
u Uniqueness -> Uniqueness -> Bool
forall a. Eq a => a -> a -> Bool
== Uniqueness
Unique
        }
    inputManifest :: ((Name, Uniqueness), ExternalValue) -> Input
inputManifest ((Name
v, Uniqueness
u), ExternalValue
vd) =
      Manifest.Input
        { inputName :: Text
Manifest.inputName = Name -> Text
nameToText Name
v,
          inputType :: Text
Manifest.inputType = ExternalValue -> Text
vdType ExternalValue
vd,
          inputUnique :: Bool
Manifest.inputUnique = Uniqueness
u Uniqueness -> Uniqueness -> Bool
forall a. Eq a => a -> a -> Bool
== Uniqueness
Unique
        }