{-# LANGUAGE QuasiQuotes #-}

-- | Code generation for server executables.
module Futhark.CodeGen.Backends.GenericC.Server
  ( serverDefs,
  )
where

import Data.Bifunctor (first, second)
import Data.Map qualified as M
import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericC.Options
import Futhark.CodeGen.Backends.GenericC.Pretty
import Futhark.CodeGen.Backends.SimpleRep
import Futhark.CodeGen.RTS.C (serverH, tuningH, valuesH)
import Futhark.Manifest
import Futhark.Util (zEncodeString)
import Language.C.Quote.OpenCL qualified as C
import Language.C.Syntax qualified as C
import Language.Futhark.Core (nameFromText)

genericOptions :: [Option]
genericOptions :: [Option]
genericOptions =
  [ Option
      { optionLongName :: String
optionLongName = String
"debugging",
        optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
'D',
        optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
        optionDescription :: String
optionDescription = String
"Perform possibly expensive internal correctness checks and verbose logging.",
        optionAction :: Stm
optionAction = [C.cstm|futhark_context_config_set_debugging(cfg, 1);|]
      },
    Option
      { optionLongName :: String
optionLongName = String
"log",
        optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
'L',
        optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
        optionDescription :: String
optionDescription = String
"Print various low-overhead logging information while running.",
        optionAction :: Stm
optionAction = [C.cstm|futhark_context_config_set_logging(cfg, 1);|]
      },
    Option
      { optionLongName :: String
optionLongName = String
"help",
        optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
'h',
        optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
        optionDescription :: String
optionDescription = String
"Print help information and exit.",
        optionAction :: Stm
optionAction =
          [C.cstm|{
                   printf("Usage: %s [OPTIONS]...\nOptions:\n\n%s\nFor more information, consult the Futhark User's Guide or the man pages.\n",
                          fut_progname, option_descriptions);
                   exit(0);
                  }|]
      },
    Option
      { optionLongName :: String
optionLongName = String
"print-params",
        optionShortName :: Maybe Char
optionShortName = forall a. Maybe a
Nothing,
        optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
        optionDescription :: String
optionDescription = String
"Print all tuning parameters that can be set with --param or --tuning.",
        optionAction :: Stm
optionAction =
          [C.cstm|{
                int n = futhark_get_tuning_param_count();
                for (int i = 0; i < n; i++) {
                  printf("%s (%s)\n", futhark_get_tuning_param_name(i),
                                      futhark_get_tuning_param_class(i));
                }
                exit(0);
              }|]
      },
    Option
      { optionLongName :: String
optionLongName = String
"param",
        optionShortName :: Maybe Char
optionShortName = forall a. Maybe a
Nothing,
        optionArgument :: OptionArgument
optionArgument = String -> OptionArgument
RequiredArgument String
"ASSIGNMENT",
        optionDescription :: String
optionDescription = String
"Set a tuning parameter to the given value.",
        optionAction :: Stm
optionAction =
          [C.cstm|{
                char *name = optarg;
                char *equals = strstr(optarg, "=");
                char *value_str = equals != NULL ? equals+1 : optarg;
                int value = atoi(value_str);
                if (equals != NULL) {
                  *equals = 0;
                  if (futhark_context_config_set_tuning_param(cfg, name, value) != 0) {
                    futhark_panic(1, "Unknown size: %s\n", name);
                  }
                } else {
                  futhark_panic(1, "Invalid argument for size option: %s\n", optarg);
                }}|]
      },
    Option
      { optionLongName :: String
optionLongName = String
"tuning",
        optionShortName :: Maybe Char
optionShortName = forall a. Maybe a
Nothing,
        optionArgument :: OptionArgument
optionArgument = String -> OptionArgument
RequiredArgument String
"FILE",
        optionDescription :: String
optionDescription = String
"Read size=value assignments from the given file.",
        optionAction :: Stm
optionAction =
          [C.cstm|{
                char *ret = load_tuning_file(optarg, cfg, (int(*)(void*, const char*, size_t))
                                                          futhark_context_config_set_tuning_param);
                if (ret != NULL) {
                  futhark_panic(1, "When loading tuning from '%s': %s\n", optarg, ret);
                }}|]
      },
    Option
      { optionLongName :: String
optionLongName = String
"cache-file",
        optionShortName :: Maybe Char
optionShortName = forall a. Maybe a
Nothing,
        optionArgument :: OptionArgument
optionArgument = String -> OptionArgument
RequiredArgument String
"FILE",
        optionDescription :: String
optionDescription = String
"Store program cache here.",
        optionAction :: Stm
optionAction =
          [C.cstm|futhark_context_config_set_cache_file(cfg, optarg);|]
      }
  ]

typeStructName :: T.Text -> String
typeStructName :: TypeName -> String
typeStructName TypeName
tname = String
"type_" forall a. Semigroup a => a -> a -> a
<> String -> String
zEncodeString (TypeName -> String
T.unpack TypeName
tname)

cType :: Manifest -> TypeName -> C.Type
cType :: Manifest -> TypeName -> Type
cType Manifest
manifest TypeName
tname =
  case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeName
tname forall a b. (a -> b) -> a -> b
$ Manifest -> Map TypeName Type
manifestTypes Manifest
manifest of
    Just (TypeArray TypeName
ctype TypeName
_ Int
_ ArrayOps
_) -> [C.cty|typename $id:(T.unpack ctype)|]
    Just (TypeOpaque TypeName
ctype OpaqueOps
_ Maybe RecordOps
_) -> [C.cty|typename $id:(T.unpack ctype)|]
    Maybe Type
Nothing -> forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Signedness -> PrimType -> Type
primAPIType forall a b. (a -> b) -> a -> b
$ TypeName -> (Signedness, PrimType)
scalarToPrim TypeName
tname

-- First component is forward declaration so we don't have to worry
-- about ordering.
typeBoilerplate :: Manifest -> (T.Text, Type) -> (C.Definition, C.Initializer, [C.Definition])
typeBoilerplate :: Manifest
-> (TypeName, Type) -> (Definition, Initializer, [Definition])
typeBoilerplate Manifest
_ (TypeName
tname, TypeArray TypeName
_ TypeName
et Int
rank ArrayOps
ops) =
  let type_name :: String
type_name = TypeName -> String
typeStructName TypeName
tname
      aux_name :: String
aux_name = String
type_name forall a. [a] -> [a] -> [a]
++ String
"_aux"
      info_name :: String
info_name = TypeName -> String
T.unpack TypeName
et forall a. [a] -> [a] -> [a]
++ String
"_info"
      shape_args :: [Exp]
shape_args = [[C.cexp|shape[$int:i]|] | Int
i <- [Int
0 .. Int
rank forall a. Num a => a -> a -> a
- Int
1]]
      array_new_wrap :: TypeName
array_new_wrap = ArrayOps -> TypeName
arrayNew ArrayOps
ops forall a. Semigroup a => a -> a -> a
<> TypeName
"_wrap"
   in ( [C.cedecl|const struct type $id:type_name;|],
        [C.cinit|&$id:type_name|],
        [C.cunit|
              void* $id:array_new_wrap(struct futhark_context *ctx,
                                       const void* p,
                                       const typename int64_t* shape) {
                return $id:(arrayNew ops)(ctx, p, $args:shape_args);
              }
              const struct array_aux $id:aux_name = {
                .name = $string:(T.unpack tname),
                .rank = $int:rank,
                .info = &$id:info_name,
                .new = (typename array_new_fn)$id:array_new_wrap,
                .free = (typename array_free_fn)$id:(arrayFree ops),
                .shape = (typename array_shape_fn)$id:(arrayShape ops),
                .values = (typename array_values_fn)$id:(arrayValues ops)
              };
              const struct type $id:type_name = {
                .name = $string:(T.unpack tname),
                .restore = (typename restore_fn)restore_array,
                .store = (typename store_fn)store_array,
                .free = (typename free_fn)free_array,
                .aux = &$id:aux_name
              };|]
      )
typeBoilerplate Manifest
manifest (TypeName
tname, TypeOpaque TypeName
c_type_name OpaqueOps
ops Maybe RecordOps
record) =
  let type_name :: String
type_name = TypeName -> String
typeStructName TypeName
tname
      aux_name :: String
aux_name = String
type_name forall a. Semigroup a => a -> a -> a
<> String
"_aux"
      ([Definition]
record_edecls, Initializer
record_init) = forall {a}.
(ToIdent a, Semigroup a, IsString a) =>
a -> Maybe RecordOps -> ([Definition], Initializer)
recordDefs String
type_name Maybe RecordOps
record
   in ( [C.cedecl|const struct type $id:type_name;|],
        [C.cinit|&$id:type_name|],
        [Definition]
record_edecls
          forall a. [a] -> [a] -> [a]
++ [C.cunit|
              const struct opaque_aux $id:aux_name = {
                .store = (typename opaque_store_fn)$id:(opaqueStore ops),
                .restore = (typename opaque_restore_fn)$id:(opaqueRestore ops),
                .free = (typename opaque_free_fn)$id:(opaqueFree ops)
              };
              const struct type $id:type_name = {
                .name = $string:(T.unpack tname),
                .restore = (typename restore_fn)restore_opaque,
                .store = (typename store_fn)store_opaque,
                .free = (typename free_fn)free_opaque,
                .aux = &$id:aux_name,
                .record = $init:record_init
              };|]
      )
  where
    recordDefs :: a -> Maybe RecordOps -> ([Definition], Initializer)
recordDefs a
_ Maybe RecordOps
Nothing = ([], [C.cinit|NULL|])
    recordDefs a
type_name (Just (RecordOps [RecordField]
fields TypeName
new)) =
      let new_wrap :: TypeName
new_wrap = TypeName
new forall a. Semigroup a => a -> a -> a
<> TypeName
"_wrap"
          record_name :: a
record_name = a
type_name forall a. Semigroup a => a -> a -> a
<> a
"_record"
          fields_name :: a
fields_name = a
type_name forall a. Semigroup a => a -> a -> a
<> a
"_fields"
          onField :: Int -> RecordField -> (Initializer, BlockItem, Exp)
onField Int
i (RecordField TypeName
name TypeName
field_tname TypeName
project) =
            let field_c_type :: Type
field_c_type = Manifest -> TypeName -> Type
cType Manifest
manifest TypeName
field_tname
                field_v :: String
field_v = String
"v" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (Int
i :: Int)
             in ( [C.cinit|{.name = $string:(T.unpack name),
                            .type = &$id:(typeStructName field_tname),
                            .project = (typename project_fn)$id:project
                           }|],
                  [C.citem|const $ty:field_c_type $id:field_v =
                            *(const $ty:field_c_type*)fields[$int:i];|],
                  [C.cexp|$id:field_v|]
                )
          ([Initializer]
field_inits, [BlockItem]
get_fields, [Exp]
field_args) = forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> RecordField -> (Initializer, BlockItem, Exp)
onField [Int
0 ..] [RecordField]
fields
       in ( [C.cunit|
             const struct field $id:fields_name[] = {
               $inits:field_inits
             };
             int $id:new_wrap(struct futhark_context* ctx, void** outp, const void* fields[]) {
               typename $id:c_type_name *out = (typename $id:c_type_name*) outp;
               $items:get_fields
               return $id:new(ctx, out, $args:field_args);
             }
             const struct record $id:record_name = {
               .num_fields = $int:(length fields),
               .fields = $id:fields_name,
               .new = $id:new_wrap
             };|],
            [C.cinit|&$id:record_name|]
          )

entryTypeBoilerplate :: Manifest -> ([C.Definition], [C.Initializer], [C.Definition])
entryTypeBoilerplate :: Manifest -> ([Definition], [Initializer], [Definition])
entryTypeBoilerplate Manifest
manifest =
  forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (Manifest
-> (TypeName, Type) -> (Definition, Initializer, [Definition])
typeBoilerplate Manifest
manifest) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. Manifest -> Map TypeName Type
manifestTypes forall a b. (a -> b) -> a -> b
$
    Manifest
manifest

oneEntryBoilerplate :: Manifest -> (T.Text, EntryPoint) -> ([C.Definition], C.Initializer)
oneEntryBoilerplate :: Manifest -> (TypeName, EntryPoint) -> ([Definition], Initializer)
oneEntryBoilerplate Manifest
manifest (TypeName
name, EntryPoint TypeName
cfun [Output]
outputs [Input]
inputs) =
  let call_f :: Name
call_f = Name
"call_" forall a. Semigroup a => a -> a -> a
<> TypeName -> Name
nameFromText TypeName
name
      out_types :: [TypeName]
out_types = forall a b. (a -> b) -> [a] -> [b]
map Output -> TypeName
outputType [Output]
outputs
      in_types :: [TypeName]
in_types = forall a b. (a -> b) -> [a] -> [b]
map Input -> TypeName
inputType [Input]
inputs
      out_types_name :: Name
out_types_name = TypeName -> Name
nameFromText TypeName
name forall a. Semigroup a => a -> a -> a
<> Name
"_out_types"
      in_types_name :: Name
in_types_name = TypeName -> Name
nameFromText TypeName
name forall a. Semigroup a => a -> a -> a
<> Name
"_in_types"
      out_unique_name :: Name
out_unique_name = TypeName -> Name
nameFromText TypeName
name forall a. Semigroup a => a -> a -> a
<> Name
"_out_unique"
      in_unique_name :: Name
in_unique_name = TypeName -> Name
nameFromText TypeName
name forall a. Semigroup a => a -> a -> a
<> Name
"_in_unique"
      ([BlockItem]
out_items, [Exp]
out_args)
        | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeName]
out_types = ([C.citems|(void)outs;|], forall a. Monoid a => a
mempty)
        | Bool
otherwise = forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> TypeName -> (BlockItem, Exp)
loadOut [Int
0 ..] [TypeName]
out_types
      ([BlockItem]
in_items, [Exp]
in_args)
        | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeName]
in_types = ([C.citems|(void)ins;|], forall a. Monoid a => a
mempty)
        | Bool
otherwise = forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> TypeName -> (BlockItem, Exp)
loadIn [Int
0 ..] [TypeName]
in_types
   in ( [C.cunit|
                const struct type* $id:out_types_name[] = {
                  $inits:(map typeStructInit out_types),
                  NULL
                };
                bool $id:out_unique_name[] = {
                  $inits:(map outputUniqueInit outputs)
                };
                const struct type* $id:in_types_name[] = {
                  $inits:(map typeStructInit in_types),
                  NULL
                };
                bool $id:in_unique_name[] = {
                  $inits:(map inputUniqueInit inputs)
                };
                int $id:call_f(struct futhark_context *ctx, void **outs, void **ins) {
                  $items:out_items
                  $items:in_items
                  return $id:cfun(ctx, $args:out_args, $args:in_args);
                }
                |],
        [C.cinit|{
            .name = $string:(T.unpack name),
            .f = $id:call_f,
            .in_types = $id:in_types_name,
            .out_types = $id:out_types_name,
            .in_unique = $id:in_unique_name,
            .out_unique = $id:out_unique_name
            }|]
      )
  where
    typeStructInit :: TypeName -> Initializer
typeStructInit TypeName
tname = [C.cinit|&$id:(typeStructName tname)|]
    inputUniqueInit :: Input -> Initializer
inputUniqueInit = Bool -> Initializer
uniqueInit forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> Bool
inputUnique
    outputUniqueInit :: Output -> Initializer
outputUniqueInit = Bool -> Initializer
uniqueInit forall b c a. (b -> c) -> (a -> b) -> a -> c
. Output -> Bool
outputUnique
    uniqueInit :: Bool -> Initializer
uniqueInit Bool
True = [C.cinit|true|]
    uniqueInit Bool
False = [C.cinit|false|]

    loadOut :: Int -> TypeName -> (BlockItem, Exp)
loadOut Int
i TypeName
tname =
      let v :: String
v = String
"out" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
i :: Int)
       in ( [C.citem|$ty:(cType manifest tname) *$id:v = outs[$int:i];|],
            [C.cexp|$id:v|]
          )
    loadIn :: Int -> TypeName -> (BlockItem, Exp)
loadIn Int
i TypeName
tname =
      let v :: String
v = String
"in" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
i :: Int)
       in ( [C.citem|$ty:(cType manifest tname) $id:v = *($ty:(cType manifest tname)*)ins[$int:i];|],
            [C.cexp|$id:v|]
          )

entryBoilerplate :: Manifest -> ([C.Definition], [C.Initializer])
entryBoilerplate :: Manifest -> ([Definition], [Initializer])
entryBoilerplate Manifest
manifest =
  forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$
    forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$
      forall a b. (a -> b) -> [a] -> [b]
map (Manifest -> (TypeName, EntryPoint) -> ([Definition], Initializer)
oneEntryBoilerplate Manifest
manifest) forall a b. (a -> b) -> a -> b
$
        forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$
          Manifest -> Map TypeName EntryPoint
manifestEntryPoints Manifest
manifest

mkBoilerplate ::
  Manifest ->
  ([C.Definition], [C.Initializer], [C.Initializer])
mkBoilerplate :: Manifest -> ([Definition], [Initializer], [Initializer])
mkBoilerplate Manifest
manifest =
  let ([Definition]
type_decls, [Initializer]
type_inits, [Definition]
type_defs) = Manifest -> ([Definition], [Initializer], [Definition])
entryTypeBoilerplate Manifest
manifest
      ([Definition]
entry_defs, [Initializer]
entry_inits) = Manifest -> ([Definition], [Initializer])
entryBoilerplate Manifest
manifest
      scalar_type_inits :: [Initializer]
scalar_type_inits = forall a b. (a -> b) -> [a] -> [b]
map TypeName -> Initializer
scalarTypeInit [TypeName]
scalar_types
   in ([Definition]
type_decls forall a. [a] -> [a] -> [a]
++ [Definition]
type_defs forall a. [a] -> [a] -> [a]
++ [Definition]
entry_defs, [Initializer]
scalar_type_inits forall a. [a] -> [a] -> [a]
++ [Initializer]
type_inits, [Initializer]
entry_inits)
  where
    scalarTypeInit :: TypeName -> Initializer
scalarTypeInit TypeName
tname = [C.cinit|&$id:(typeStructName tname)|]
    scalar_types :: [TypeName]
scalar_types =
      [ TypeName
"i8",
        TypeName
"i16",
        TypeName
"i32",
        TypeName
"i64",
        TypeName
"u8",
        TypeName
"u16",
        TypeName
"u32",
        TypeName
"u64",
        TypeName
"f16",
        TypeName
"f32",
        TypeName
"f64",
        TypeName
"bool"
      ]

{-# NOINLINE serverDefs #-}

-- | Generate Futhark server executable code.
serverDefs :: [Option] -> Manifest -> T.Text
serverDefs :: [Option] -> Manifest -> TypeName
serverDefs [Option]
options Manifest
manifest =
  let option_parser :: Func
option_parser =
        String -> [Option] -> Func
generateOptionParser String
"parse_options" forall a b. (a -> b) -> a -> b
$ [Option]
genericOptions forall a. [a] -> [a] -> [a]
++ [Option]
options
      ([Definition]
boilerplate_defs, [Initializer]
type_inits, [Initializer]
entry_point_inits) =
        Manifest -> ([Definition], [Initializer], [Initializer])
mkBoilerplate Manifest
manifest
   in [Definition] -> TypeName
definitionsText
        [C.cunit|
$esc:("#include <getopt.h>")
$esc:("#include <ctype.h>")
$esc:("#include <inttypes.h>")

// If the entry point is NULL, the program will terminate after doing initialisation and such.  It is not used for anything else in server mode.
static const char *entry_point = "main";

$esc:(T.unpack valuesH)
$esc:(T.unpack serverH)
$esc:(T.unpack tuningH)

$edecls:boilerplate_defs

const struct type* types[] = {
  $inits:type_inits,
  NULL
};

struct entry_point entry_points[] = {
  $inits:entry_point_inits,
  { .name = NULL }
};

struct futhark_prog prog = {
  .types = types,
  .entry_points = entry_points
};

$func:option_parser

int main(int argc, char** argv) {
  fut_progname = argv[0];

  struct futhark_context_config *cfg = futhark_context_config_new();
  assert(cfg != NULL);

  int parsed_options = parse_options(cfg, argc, argv);
  argc -= parsed_options;
  argv += parsed_options;

  if (argc != 0) {
    futhark_panic(1, "Excess non-option: %s\n", argv[0]);
  }

  struct futhark_context *ctx = futhark_context_new(cfg);
  assert (ctx != NULL);

  futhark_context_set_logging_file(ctx, stdout);

  char* error = futhark_context_get_error(ctx);
  if (error != NULL) {
    futhark_panic(1, "Error during context initialisation:\n%s", error);
  }

  if (entry_point != NULL) {
    run_server(&prog, cfg, ctx);
  }

  futhark_context_free(ctx);
  futhark_context_config_free(cfg);
}
|]