content: {-# LANGUAGE QuasiQuotes #-}
  
  -- | C code generator.  This module can convert a correct ImpCode
  -- program to an equivalent C program.
  module Futhark.CodeGen.Backends.MulticoreC
    ( compileProg,
      GC.CParts (..),
      GC.asLibrary,
      GC.asExecutable,
      GC.asServer,
      operations,
      cliOptions,
      compileOp,
      ValueType (..),
      paramToCType,
      prepareTaskStruct,
      closureFreeStructField,
      generateParLoopFn,
      addTimingFields,
      functionTiming,
      functionIterations,
      multicoreDef,
      multicoreName,
      DefSpecifier,
      atomicOps,
    )
  where
  
  import Control.Monad
  import Data.Loc
  import Data.Map qualified as M
  import Data.Maybe
  import Data.Text qualified as T
  import Futhark.CodeGen.Backends.GenericC qualified as GC
  import Futhark.CodeGen.Backends.GenericC.Options
  import Futhark.CodeGen.Backends.MulticoreC.Boilerplate (generateBoilerplate)
  import Futhark.CodeGen.Backends.SimpleRep
  import Futhark.CodeGen.ImpCode.Multicore hiding (ValueType)
  import Futhark.CodeGen.ImpGen.Multicore qualified as ImpGen
  import Futhark.IR.MCMem (MCMem, Prog)
  import Futhark.MonadFreshNames
  import Language.C.Quote.OpenCL qualified as C
  import Language.C.Syntax qualified as C
  
  -- | Compile the program to ImpCode with multicore operations.
  compileProg ::
    (MonadFreshNames m) => T.Text -> Prog MCMem -> m (ImpGen.Warnings, GC.CParts)
  compileProg version =
    traverse
      ( GC.compileProg
          "multicore"
          version
          mempty
          operations
          generateBoilerplate
          "#include <pthread.h>\n"
          (DefaultSpace, [DefaultSpace])
          cliOptions
      )
      <=< ImpGen.compileProg
  
  -- | Multicore-related command line options.
  cliOptions :: [Option]
  cliOptions =
    [ Option
        { optionLongName = "num-threads",
          optionShortName = Nothing,
          optionArgument = RequiredArgument "INT",
          optionAction = [C.cstm|futhark_context_config_set_num_threads(cfg, atoi(optarg));|],
          optionDescription = "Set number of threads used for execution."
        }
    ]
  
  -- | Operations for generating multicore code.
  operations :: GC.Operations Multicore s
  operations =
    GC.defaultOperations
      { GC.opsCompiler = compileOp,
        GC.opsCritical =
          -- The thread entering an API function is always considered
          -- the "first worker" - note that this might differ from the
          -- thread that created the context!  This likely only matters
          -- for entry points, since they are the only API functions
          -- that contain parallel operations.
          ( [C.citems|worker_local = &ctx->scheduler.workers[0];|],
            []
          )
      }
  
  closureFreeStructField :: VName -> Name
  closureFreeStructField v =
    nameFromString "free_" <> nameFromString (prettyString v)
  
  closureRetvalStructField :: VName -> Name
  closureRetvalStructField v =
    nameFromString "retval_" <> nameFromString (prettyString v)
  
  data ValueType = Prim PrimType | MemBlock | RawMem
  
  compileFreeStructFields :: [VName] -> [(C.Type, ValueType)] -> [C.FieldGroup]
  compileFreeStructFields = zipWith field
    where
      field name (ty, Prim _) =
        [C.csdecl|$ty:ty $id:(closureFreeStructField name);|]
      field name (_, _) =
        [C.csdecl|$ty:defaultMemBlockType $id:(closureFreeStructField name);|]
  
  compileRetvalStructFields :: [VName] -> [(C.Type, ValueType)] -> [C.FieldGroup]
  compileRetvalStructFields = zipWith field
    where
      field name (ty, Prim _) =
        [C.csdecl|$ty:ty *$id:(closureRetvalStructField name);|]
      field name (_, _) =
        [C.csdecl|$ty:defaultMemBlockType $id:(closureRetvalStructField name);|]
  
  compileSetStructValues ::
    (C.ToIdent a) =>
    a ->
    [VName] ->
    [(C.Type, ValueType)] ->
    [C.Stm]
  compileSetStructValues struct = zipWith field
    where
      field name (_, Prim pt) =
        [C.cstm|$id:struct.$id:(closureFreeStructField name)=$exp:(toStorage pt (C.toExp name noLoc));|]
      field name (_, MemBlock) =
        [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name.mem;|]
      field name (_, RawMem) =
        [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name;|]
  
  compileSetRetvalStructValues ::
    (C.ToIdent a) =>
    a ->
    [VName] ->
    [(C.Type, ValueType)] ->
    [C.Stm]
  compileSetRetvalStructValues struct vnames we = concat $ zipWith field vnames we
    where
      field name (ct, Prim _) =
        [C.cstms|$id:struct.$id:(closureRetvalStructField name)=(($ty:ct*)&$id:name);
                 $escstm:("#if defined(ISPC)")
                 $id:struct.$id:(closureRetvalStructField name)+= programIndex;
                 $escstm:("#endif")|]
      field name (_, MemBlock) =
        [C.cstms|$id:struct.$id:(closureRetvalStructField name)=$id:name.mem;|]
      field name (_, RawMem) =
        [C.cstms|$id:struct.$id:(closureRetvalStructField name)=$id:name;|]
  
  compileGetRetvalStructVals :: (C.ToIdent a) => a -> [VName] -> [(C.Type, ValueType)] -> [C.InitGroup]
  compileGetRetvalStructVals struct = zipWith field
    where
      field name (ty, Prim pt) =
        let inner = [C.cexp|*$id:struct->$id:(closureRetvalStructField name)|]
         in [C.cdecl|$ty:ty $id:name = $exp:(fromStorage pt inner);|]
      field name (ty, _) =
        [C.cdecl|$ty:ty $id:name =
                   {.desc = $string:(prettyString name),
                   .mem = $id:struct->$id:(closureRetvalStructField name),
                   .size = 0, .references = NULL};|]
  
  compileGetStructVals ::
    (C.ToIdent a) =>
    a ->
    [VName] ->
    [(C.Type, ValueType)] ->
    [C.InitGroup]
  compileGetStructVals struct = zipWith field
    where
      field name (ty, Prim pt) =
        let inner = [C.cexp|$id:struct->$id:(closureFreeStructField name)|]
         in [C.cdecl|$ty:ty $id:name = $exp:(fromStorage pt inner);|]
      field name (ty, _) =
        [C.cdecl|$ty:ty $id:name =
                   {.desc = $string:(prettyString name),
                    .mem = $id:struct->$id:(closureFreeStructField name),
                    .size = 0, .references = NULL};|]
  
  compileWriteBackResVals :: (C.ToIdent a) => a -> [VName] -> [(C.Type, ValueType)] -> [C.Stm]
  compileWriteBackResVals struct = zipWith field
    where
      field name (_, Prim pt) =
        [C.cstm|*$id:struct->$id:(closureRetvalStructField name) = $exp:(toStorage pt (C.toExp name noLoc));|]
      field name (_, _) =
        [C.cstm|$id:struct->$id:(closureRetvalStructField name) = $id:name.mem;|]
  
  paramToCType :: Param -> GC.CompilerM op s (C.Type, ValueType)
  paramToCType (ScalarParam _ pt) = do
    let t = primStorageType pt
    pure (t, Prim pt)
  paramToCType (MemParam name space') = mcMemToCType name space'
  
  mcMemToCType :: VName -> Space -> GC.CompilerM op s (C.Type, ValueType)
  mcMemToCType v space = do
    refcount <- GC.fatMemory space
    cached <- isJust <$> GC.cacheMem v
    pure
      ( GC.fatMemType space,
        if refcount && not cached
          then MemBlock
          else RawMem
      )
  
  benchmarkCode :: Name -> [C.BlockItem] -> GC.CompilerM op s [C.BlockItem]
  benchmarkCode name code = do
    event <- newVName "event"
    provenance <- GC.provenanceExp
    pure
      [C.citems|
       struct mc_event* $id:event = mc_event_new(ctx);
       if ($id:event != NULL) {
         $id:event->bef = get_wall_time();
       }
       $items:code
       if ($id:event != NULL) {
         $id:event->aft = get_wall_time();
         lock_lock(&ctx->event_list_lock);
         add_event(ctx,
                   $string:(nameToString name),
                   $exp:provenance,
                   NULL,
                   $id:event,
                   (typename event_report_fn)mc_event_report);
         lock_unlock(&ctx->event_list_lock);
       }|]
  
  functionTiming :: Name -> C.Id
  functionTiming = (`C.toIdent` mempty) . (<> "_total_time")
  
  functionIterations :: Name -> C.Id
  functionIterations = (`C.toIdent` mempty) . (<> "_total_iter")
  
  addTimingFields :: Name -> GC.CompilerM op s ()
  addTimingFields name = do
    GC.contextField (functionTiming name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]
    GC.contextField (functionIterations name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]
  
  multicoreName :: Name -> GC.CompilerM op s Name
  multicoreName s = do
    s' <- newVName $ "futhark_mc_" <> s
    pure $ baseName s' <> "_" <> nameFromString (show (baseTag s'))
  
  type DefSpecifier s = Name -> (Name -> GC.CompilerM Multicore s C.Definition) -> GC.CompilerM Multicore s Name
  
  multicoreDef :: DefSpecifier s
  multicoreDef s f = do
    s' <- multicoreName s
    GC.libDecl =<< f s'
    pure s'
  
  generateParLoopFn ::
    (C.ToIdent a) =>
    M.Map VName Space ->
    Name ->
    MCCode ->
    a ->
    [(VName, (C.Type, ValueType))] ->
    [(VName, (C.Type, ValueType))] ->
    GC.CompilerM Multicore s Name
  generateParLoopFn lexical basename code fstruct free retval = do
    let (fargs, fctypes) = unzip free
    let (retval_args, retval_ctypes) = unzip retval
    multicoreDef basename $ \s -> do
      fbody <- benchmarkCode s <=< GC.inNewFunction $
        GC.cachingMemory lexical $ \decl_cached free_cached -> GC.collect $ do
          mapM_ GC.item [C.citems|$decls:(compileGetStructVals fstruct fargs fctypes)|]
          mapM_ GC.item [C.citems|$decls:(compileGetRetvalStructVals fstruct retval_args retval_ctypes)|]
          code' <- GC.collect $ GC.compileCode code
          mapM_ GC.item decl_cached
          mapM_ GC.item =<< GC.declAllocatedMem
          mapM_ GC.item code'
          free_mem <- GC.freeAllocatedMem
          GC.stm [C.cstm|cleanup: {$stms:free_cached $items:free_mem}|]
      pure
        [C.cedecl|int $id:s(void *args, typename int64_t iterations, int tid, struct scheduler_info info) {
                             int err = 0;
                             int subtask_id = tid;
                             struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args;
                             struct futhark_context *ctx = $id:fstruct->ctx;
                             $items:fbody
                             if (err == 0) {
                               $stms:(compileWriteBackResVals fstruct retval_args retval_ctypes)
                             }
                             return err;
                        }|]
  
  prepareTaskStruct ::
    DefSpecifier s ->
    Name ->
    [VName] ->
    [(C.Type, ValueType)] ->
    [VName] ->
    [(C.Type, ValueType)] ->
    GC.CompilerM Multicore s Name
  prepareTaskStruct def name free_args free_ctypes retval_args retval_ctypes = do
    let makeStruct s =
          pure
            [C.cedecl|struct $id:s {
                         struct futhark_context *ctx;
                         $sdecls:(compileFreeStructFields free_args free_ctypes)
                         $sdecls:(compileRetvalStructFields retval_args retval_ctypes)
                       };|]
    fstruct <- def name makeStruct
    let fstruct' = fstruct <> "_"
    GC.decl [C.cdecl|struct $id:fstruct $id:fstruct';|]
    GC.stm [C.cstm|$id:fstruct'.ctx = ctx;|]
    GC.stms [C.cstms|$stms:(compileSetStructValues fstruct' free_args free_ctypes)|]
    GC.stms [C.cstms|$stms:(compileSetRetvalStructValues fstruct' retval_args retval_ctypes)|]
    pure fstruct
  
  -- Generate a segop function for top_level and potentially nested SegOp code
  compileOp :: GC.OpCompiler Multicore s
  compileOp (GetLoopBounds start end) = do
    GC.stm [C.cstm|$id:start = start;|]
    GC.stm [C.cstm|$id:end = end;|]
  compileOp (GetTaskId v) =
    GC.stm [C.cstm|$id:v = subtask_id;|]
  compileOp (GetNumTasks v) =
    GC.stm [C.cstm|$id:v = info.nsubtasks;|]
  compileOp (SegOp name params seq_task par_task retvals (SchedulerInfo e sched)) = do
    let (ParallelTask seq_code) = seq_task
    free_ctypes <- mapM paramToCType params
    retval_ctypes <- mapM paramToCType retvals
    let free_args = map paramName params
        retval_args = map paramName retvals
        free = zip free_args free_ctypes
        retval = zip retval_args retval_ctypes
  
    e' <- GC.compileExp e
  
    let lexical = lexicalMemoryUsageMC TraverseKernels $ Function Nothing [] params seq_code
  
    fstruct <-
      prepareTaskStruct multicoreDef "task" free_args free_ctypes retval_args retval_ctypes
  
    fpar_task <- generateParLoopFn lexical (name <> "_task") seq_code fstruct free retval
    addTimingFields fpar_task
  
    let ftask_name = fstruct <> "_task"
    GC.decl [C.cdecl|struct scheduler_segop $id:ftask_name;|]
    GC.stm [C.cstm|$id:ftask_name.args = &$id:(fstruct <> "_");|]
    GC.stm [C.cstm|$id:ftask_name.top_level_fn = $id:fpar_task;|]
    GC.stm [C.cstm|$id:ftask_name.name = $string:(nameToString fpar_task);|]
    GC.stm [C.cstm|$id:ftask_name.iterations = $exp:e';|]
    -- Create the timing fields for the task
    GC.stm [C.cstm|$id:ftask_name.task_time = &ctx->program->$id:(functionTiming fpar_task);|]
    GC.stm [C.cstm|$id:ftask_name.task_iter = &ctx->program->$id:(functionIterations fpar_task);|]
  
    case sched of
      Dynamic -> GC.stm [C.cstm|$id:ftask_name.sched = DYNAMIC;|]
      Static -> GC.stm [C.cstm|$id:ftask_name.sched = STATIC;|]
  
    -- Generate the nested segop function if available
    case par_task of
      Just (ParallelTask nested_code) -> do
        let lexical_nested = lexicalMemoryUsageMC TraverseKernels $ Function Nothing [] params nested_code
        fnpar_task <- generateParLoopFn lexical_nested (name <> "_nested_task") nested_code fstruct free retval
        GC.stm [C.cstm|$id:ftask_name.nested_fn = $id:fnpar_task;|]
      Nothing ->
        GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|]
  
    let ftask_err = fpar_task <> "_err"
        code =
          [C.citems|int $id:ftask_err = scheduler_prepare_task(&ctx->scheduler, &$id:ftask_name);
                    if ($id:ftask_err != 0) {
                      err = $id:ftask_err;
                      goto cleanup;
                    }|]
  
    mapM_ GC.item code
  compileOp (ParLoop s' body free) = do
    free_ctypes <- mapM paramToCType free
    let free_args = map paramName free
  
    let lexical = lexicalMemoryUsageMC TraverseKernels $ Function Nothing [] free body
  
    fstruct <-
      prepareTaskStruct multicoreDef (s' <> "_parloop_struct") free_args free_ctypes mempty mempty
  
    ftask <- multicoreDef (s' <> "_parloop") $ \s -> do
      fbody <- benchmarkCode s <=< GC.inNewFunction $
        GC.cachingMemory lexical $ \decl_cached free_cached -> GC.collect $ do
          GC.items [C.citems|$decls:(compileGetStructVals fstruct free_args free_ctypes)|]
  
          GC.decl [C.cdecl|typename int64_t iterations = end-start;|]
  
          body' <- GC.collect $ GC.compileCode body
  
          mapM_ GC.item decl_cached
          mapM_ GC.item =<< GC.declAllocatedMem
          free_mem <- GC.freeAllocatedMem
          mapM_ GC.item body'
          GC.stm [C.cstm|cleanup: {$stms:free_cached $items:free_mem}|]
      pure
        [C.cedecl|static int $id:s(void *args,
                                   typename int64_t start,
                                   typename int64_t end,
                                   int subtask_id,
                                   int tid) {
                         (void)subtask_id;
                         (void)tid;
                         int err = 0;
                         struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args;
                         struct futhark_context *ctx = $id:fstruct->ctx;
                         $items:fbody
                         return err;
                  }|]
  
    let ftask_name = ftask <> "_task"
    GC.decl [C.cdecl|struct scheduler_parloop $id:ftask_name;|]
    GC.stm [C.cstm|$id:ftask_name.name = $string:(nameToString ftask);|]
    GC.stm [C.cstm|$id:ftask_name.fn = $id:ftask;|]
    GC.stm [C.cstm|$id:ftask_name.args = &$id:(fstruct <> "_");|]
    GC.stm [C.cstm|$id:ftask_name.iterations = iterations;|]
    GC.stm [C.cstm|$id:ftask_name.info = info;|]
  
    let ftask_err = ftask <> "_err"
        ftask_total = ftask <> "_total"
    code' <-
      benchmarkCode
        ftask_total
        [C.citems|int $id:ftask_err = scheduler_execute_task(&ctx->scheduler,
                                                             &$id:ftask_name);
                 if ($id:ftask_err != 0) {
                   err = $id:ftask_err;
                   goto cleanup;
                 }|]
  
    mapM_ GC.item code'
  compileOp (Atomic aop) =
    atomicOps aop (\ty _ -> pure [C.cty|$ty:ty*|])
  compileOp (ISPCKernel body _) =
    scopedBlock body
  compileOp (ForEach i from bound body) = do
    let i' = C.toIdent i
        t = primTypeToCType $ primExpType bound
    from' <- GC.compileExp from
    bound' <- GC.compileExp bound
    body' <- GC.collect $ GC.compileCode body
    GC.stm
      [C.cstm|for ($ty:t $id:i' = $exp:from'; $id:i' < $exp:bound'; $id:i'++) {
              $items:body'
            }|]
  compileOp (ForEachActive i body) = do
    GC.decl [C.cdecl|typename int64_t $id:i = 0;|]
    scopedBlock body
  compileOp (ExtractLane dest tar _) = do
    tar' <- GC.compileExp tar
    GC.stm [C.cstm|$id:dest = $exp:tar';|]
  
  scopedBlock :: MCCode -> GC.CompilerM Multicore s ()
  scopedBlock code = do
    inner <- GC.collect $ GC.compileCode code
    GC.stm [C.cstm|{$items:inner}|]
  
  doAtomic ::
    (C.ToIdent a1) =>
    a1 ->
    VName ->
    Count u (TExp Int32) ->
    Exp ->
    String ->
    C.Type ->
    (C.Type -> VName -> GC.CompilerM op s C.Type) ->
    GC.CompilerM op s ()
  doAtomic old arr ind val op ty castf = do
    ind' <- GC.compileExp $ untyped $ unCount ind
    val' <- GC.compileExp val
    cast <- castf ty arr
    arr' <- GC.rawMem arr
    GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast)$exp:arr')[$exp:ind'], ($ty:ty) $exp:val', __ATOMIC_RELAXED);|]
  
  atomicOps :: AtomicOp -> (C.Type -> VName -> GC.CompilerM op s C.Type) -> GC.CompilerM op s ()
  atomicOps (AtomicCmpXchg t old arr ind res val) castf = do
    ind' <- GC.compileExp $ untyped $ unCount ind
    new_val' <- GC.compileExp val
    cast <- castf [C.cty|$ty:(GC.primTypeToCType t)|] arr
    arr' <- GC.rawMem arr
    GC.stm
      [C.cstm|$id:res = $id:op(&(($ty:cast)$exp:arr')[$exp:ind'],
                   &$id:old,
                   $exp:new_val',
                   0, __ATOMIC_SEQ_CST, __ATOMIC_RELAXED);|]
    where
      op :: String
      op = "__atomic_compare_exchange_n"
  atomicOps (AtomicXchg t old arr ind val) castf = do
    ind' <- GC.compileExp $ untyped $ unCount ind
    val' <- GC.compileExp val
    cast <- castf [C.cty|$ty:(GC.primTypeToCType t)|] arr
    GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast)$id:arr.mem)[$exp:ind'], $exp:val', __ATOMIC_SEQ_CST);|]
    where
      op :: String
      op = "__atomic_exchange_n"
  atomicOps (AtomicAdd t old arr ind val) castf =
    doAtomic old arr ind val "__atomic_fetch_add" [C.cty|$ty:(GC.intTypeToCType t)|] castf
  atomicOps (AtomicSub t old arr ind val) castf =
    doAtomic old arr ind val "__atomic_fetch_sub" [C.cty|$ty:(GC.intTypeToCType t)|] castf
  atomicOps (AtomicAnd t old arr ind val) castf =
    doAtomic old arr ind val "__atomic_fetch_and" [C.cty|$ty:(GC.intTypeToCType t)|] castf
  atomicOps (AtomicOr t old arr ind val) castf =
    doAtomic old arr ind val "__atomic_fetch_or" [C.cty|$ty:(GC.intTypeToCType t)|] castf
  atomicOps (AtomicXor t old arr ind val) castf =
    doAtomic old arr ind val "__atomic_fetch_xor" [C.cty|$ty:(GC.intTypeToCType t)|] castf
  