Server.hs

{-# LANGUAGE QuasiQuotes #-}

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

import Data.Bifunctor (first, second)
import Data.List (unzip4)
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 (zEncodeText)
import Futhark.Util.Pretty (prettyText)
import Language.C.Quote.OpenCL qualified as C
import Language.C.Syntax qualified as C
import Language.Futhark.Core (nameFromText)

genericOptions :: [Option]
genericOptions =
  [ Option
      { optionLongName = "debugging",
        optionShortName = Just 'D',
        optionArgument = NoArgument,
        optionDescription = "Perform possibly expensive internal correctness checks and verbose logging.",
        optionAction = [C.cstm|futhark_context_config_set_debugging(cfg, 1);|]
      },
    Option
      { optionLongName = "log",
        optionShortName = Just 'L',
        optionArgument = NoArgument,
        optionDescription = "Print various low-overhead logging information while running.",
        optionAction = [C.cstm|futhark_context_config_set_logging(cfg, 1);|]
      },
    Option
      { optionLongName = "profile",
        optionShortName = Just 'P',
        optionArgument = NoArgument,
        optionDescription = "Enable the collection of profiling information.",
        optionAction = [C.cstm|futhark_context_config_set_profiling(cfg, 1);|]
      },
    Option
      { optionLongName = "help",
        optionShortName = Just 'h',
        optionArgument = NoArgument,
        optionDescription = "Print help information and exit.",
        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 = "print-params",
        optionShortName = Nothing,
        optionArgument = NoArgument,
        optionDescription = "Print all tuning parameters that can be set with --param or --tuning.",
        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 = "param",
        optionShortName = Nothing,
        optionArgument = RequiredArgument "ASSIGNMENT",
        optionDescription = "Set a tuning parameter to the given value.",
        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 parameter: %s\n", name);
                  }
                } else {
                  futhark_panic(1, "Invalid argument for --parameter option: %s\n", optarg);
                }}|]
      },
    Option
      { optionLongName = "tuning",
        optionShortName = Nothing,
        optionArgument = RequiredArgument "FILE",
        optionDescription = "Read size=value assignments from the given file.",
        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 file '%s': %s\n", optarg, ret);
                }}|]
      },
    Option
      { optionLongName = "cache-file",
        optionShortName = Nothing,
        optionArgument = RequiredArgument "FILE",
        optionDescription = "Store program cache here.",
        optionAction =
          [C.cstm|futhark_context_config_set_cache_file(cfg, optarg);|]
      }
  ]

typeStructName :: T.Text -> T.Text
typeStructName tname = "type_" <> zEncodeText tname

cType :: Manifest -> TypeName -> C.Type
cType manifest tname =
  case M.lookup tname $ manifestTypes manifest of
    Just (TypeArray ctype _ _ _) -> [C.cty|typename $id:(T.unpack ctype)|]
    Just (TypeOpaque ctype _ _) -> [C.cty|typename $id:(T.unpack ctype)|]
    Nothing -> uncurry primAPIType $ scalarToPrim tname

data Kind
  = Primitive
  | -- TODO: Implement array functions for opaque arrays
    Array
  | Record
  | Sum
  | Opaque

cKind :: Kind -> C.Exp
cKind Primitive = [C.cexp|PRIMITIVE|]
cKind Array = [C.cexp|ARRAY|]
cKind Record = [C.cexp|RECORD|]
cKind Sum = [C.cexp|SUM|]
cKind Opaque = [C.cexp|OPAQUE|]

-- 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 _ (tname, TypeArray c_type_name et rank ops) =
  let element_type_name = typeStructName et
      type_name = typeStructName tname
      array_name = type_name <> "_array"
      aux_name = type_name <> "_aux"
      info_name = et <> "_info"
      array_new_wrap = arrayNew ops <> "_wrap"
      array_index_wrap = arrayIndex ops <> "_wrap"
      shape_args = [[C.cexp|shape[$int:i]|] | i <- [0 .. rank - 1]]
      is_args = [[C.cexp|is[$int:i]|] | i <- [0 .. rank - 1]]
   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);
              }
              int $id:array_index_wrap(struct futhark_context *ctx,
                                       void *dest,
                                       typename $id:c_type_name arr,
                                       const typename int64_t *is) {
                return $id:(arrayIndex ops)(ctx, dest, arr, $args:is_args);
              }
              const struct array $id:array_name = {
                .rank = $int:rank,
                .element_type = &$id:element_type_name,
                .new = (typename array_new_fn)$id:array_new_wrap,
                .shape = (typename array_shape_fn)$id:(arrayShape ops),
                .index = (typename array_index_fn)$id:array_index_wrap,
              };
              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,
                .kind = $exp:(cKind Array),
                .info = &$id:array_name
              };|]
      )
typeBoilerplate manifest (tname, TypeOpaque c_type_name ops extra_ops) =
  let type_name = typeStructName tname
      aux_name = type_name <> "_aux"
      (transparent_edecls, transparent_init, kind) = transparentDefs type_name extra_ops
   in ( [C.cedecl|const struct type $id:type_name;|],
        [C.cinit|&$id:type_name|],
        transparent_edecls
          ++ [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,
                .kind = $exp:(cKind kind),
                .info = $init:transparent_init
              };|]
      )
  where
    transparentDefs type_name (Just (OpaqueRecord (RecordOps fields new))) =
      let new_wrap = new <> "_wrap"
          record_name = type_name <> "_record"
          fields_name = type_name <> "_fields"
          onField i (RecordField name field_tname project) =
            let field_c_type = cType manifest field_tname
                field_v = "v" <> show (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|]
                )
          (field_inits, get_fields, field_args) = unzip3 $ zipWith onField [0 ..] 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|],
            Record
          )
    transparentDefs type_name (Just (OpaqueSum (SumOps variants variant))) =
      let sum_name = type_name <> "_sum"
          variants_name = type_name <> "_variants"
          onType i type_tname =
            let type_c_type = cType manifest type_tname
                type_v = "v" <> show (i :: Int)
             in ( [C.citem|const $ty:type_c_type $id:type_v =
                            *(const $ty:type_c_type*)types[$int:i];|],
                  [C.citem|$ty:type_c_type *$id:type_v =
                            ($ty:type_c_type*)outs[$int:i];|],
                  [C.cexp|$id:type_v|],
                  [C.cinit|&$id:(typeStructName type_tname)|]
                )
          onVariant (SumVariant name variant_tnames new destruct) =
            let new_wrap = new <> "_wrap"
                destruct_wrap = destruct <> "_wrap"
                types_name = type_name <> "_" <> name <> "_fields"
                (get_types, get_outs, type_args, type_struct_names) = unzip4 $ zipWith onType [0 ..] variant_tnames
             in ( [C.cinit|{.name = $string:(T.unpack name),
                            .num_types = $int:(length variant_tnames),
                            .types = $id:types_name,
                            .new = $id:new_wrap,
                            .destruct = $id:destruct_wrap
                           }|],
                  [C.cunit|
                    const struct type *$id:types_name[] = {
                      $inits:type_struct_names
                    };
                    int $id:new_wrap(struct futhark_context* ctx, void** outp, const void* types[]) {
                      typename $id:c_type_name *out = (typename $id:c_type_name*) outp;
                      $items:get_types
                      return $id:new(ctx, out, $args:type_args);
                    }
                    int $id:destruct_wrap(struct futhark_context* ctx, const void* outs[], const void* inp) {
                      const typename $id:c_type_name in = (const typename $id:c_type_name) inp;
                      $items:get_outs
                      return $id:destruct(ctx, $args:type_args, in);
                    }|]
                )

          (variant_inits, variant_wraps) = unzip $ map onVariant variants
       in ( foldl1 (++) variant_wraps
              ++ [C.cunit|
             const struct variant $id:variants_name[] = {
               $inits:variant_inits
             };
             const struct sum $id:sum_name = {
               .num_variants = $int:(length variants),
               .variants = $id:variants_name,
               .variant = (int (*)(struct futhark_context *, const void *))$id:variant
             };|],
            [C.cinit|&$id:sum_name|],
            Sum
          )
    transparentDefs type_name (Just (OpaqueArray ops')) = opaqueArrayDefs type_name (opaqueArrayRank ops') (opaqueArrayElemType ops') (opaqueArrayShape ops') (opaqueArrayIndex ops')
    transparentDefs type_name (Just (OpaqueRecordArray ops')) = opaqueArrayDefs type_name (recordArrayRank ops') (recordArrayElemType ops') (recordArrayShape ops') (recordArrayIndex ops')
    transparentDefs _ _ = ([], [C.cinit|NULL|], Opaque)

    opaqueArrayDefs type_name rank et shape index =
      let array_name = type_name <> "_array"
          element_type_name = typeStructName et
          index_wrap = index <> "_wrap"
          is_args = [[C.cexp|is[$int:i]|] | i <- [0 .. rank - 1]]
       in ( [C.cunit|
              int $id:index_wrap(struct futhark_context *ctx,
                                       void *dest,
                                       typename $id:c_type_name arr,
                                       const typename int64_t *is) {
                return $id:index(ctx, dest, arr, $args:is_args);
              }
              const struct array $id:array_name = {
                .rank = $int:rank,
                .element_type = &$id:element_type_name,
                .new = NULL,
                .shape = (typename array_shape_fn)$id:shape,
                .index = (typename array_index_fn)$id:index_wrap,
              };|],
            [C.cinit|&$id:array_name|],
            Array
          )

entryTypeBoilerplate :: Manifest -> ([C.Definition], [C.Initializer], [C.Definition])
entryTypeBoilerplate manifest =
  second concat . unzip3 . map (typeBoilerplate manifest) . M.toList . manifestTypes $
    manifest

oneEntryBoilerplate :: Manifest -> (T.Text, EntryPoint) -> ([C.Definition], C.Initializer)
oneEntryBoilerplate manifest (name, EntryPoint cfun tuning_params outputs inputs attrs) =
  let call_f = "call_" <> nameFromText name
      out_types = map outputType outputs
      in_types = map inputType inputs
      out_types_name = nameFromText name <> "_out_types"
      in_types_name = nameFromText name <> "_in_types"
      out_unique_name = nameFromText name <> "_out_unique"
      in_unique_name = nameFromText name <> "_in_unique"
      tuning_params_name = nameFromText name <> "_tuning_params"
      attrs_name = nameFromText name <> "_attrs"
      (out_items, out_args)
        | null out_types = ([C.citems|(void)outs;|], mempty)
        | otherwise = unzip $ zipWith loadOut [0 ..] out_types
      (in_items, in_args)
        | null in_types = ([C.citems|(void)ins;|], mempty)
        | otherwise = unzip $ zipWith loadIn [0 ..] 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)
                };
                const char* $id:tuning_params_name[] = {
                  $inits:(map textInit tuning_params),
                  NULL
                };
                const char* $id:attrs_name[] = {
                  $inits:(map (textInit . prettyText) attrs),
                  NULL
                };
                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,
            .tuning_params = $id:tuning_params_name,
            .in_types = $id:in_types_name,
            .out_types = $id:out_types_name,
            .in_unique = $id:in_unique_name,
            .out_unique = $id:out_unique_name,
            .attrs = $id:attrs_name
            }|]
      )
  where
    typeStructInit tname = [C.cinit|&$id:(typeStructName tname)|]
    inputUniqueInit = uniqueInit . inputUnique
    outputUniqueInit = uniqueInit . outputUnique
    uniqueInit True = [C.cinit|true|]
    uniqueInit False = [C.cinit|false|]

    loadOut i tname =
      let v = "out" ++ show (i :: Int)
       in ( [C.citem|$ty:(cType manifest tname) *$id:v = outs[$int:i];|],
            [C.cexp|$id:v|]
          )
    loadIn i tname =
      let v = "in" ++ show (i :: Int)
       in ( [C.citem|$ty:(cType manifest tname) $id:v = *($ty:(cType manifest tname)*)ins[$int:i];|],
            [C.cexp|$id:v|]
          )

    textInit t = [C.cinit|$string:(T.unpack t)|]

entryBoilerplate :: Manifest -> ([C.Definition], [C.Initializer])
entryBoilerplate manifest =
  first concat $
    unzip $
      map (oneEntryBoilerplate manifest) $
        M.toList $
          manifestEntryPoints manifest

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

{-# NOINLINE serverDefs #-}

-- | Generate Futhark server executable code.
serverDefs :: [Option] -> Manifest -> T.Text
serverDefs options manifest =
  let option_parser =
        generateOptionParser "parse_options" $ genericOptions ++ options
      (boilerplate_defs, type_inits, entry_point_inits) =
        mkBoilerplate manifest
   in 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);
}
|]