    {-# LANGUAGE QuasiQuotes #-}

    -- | Code generation for CUDA.
    module Futhark.CodeGen.Backends.CCUDA
      ( compileProg,
        GC.CParts (..),
        GC.asLibrary,
        GC.asExecutable,
        GC.asServer,
      )
    where

    import Data.Map qualified as M
    import Data.Text qualified as T
    import Futhark.CodeGen.Backends.GPU
    import Futhark.CodeGen.Backends.GenericC qualified as GC
    import Futhark.CodeGen.Backends.GenericC.Options
    import Futhark.CodeGen.ImpCode.OpenCL
    import Futhark.CodeGen.ImpGen.CUDA qualified as ImpGen
    import Futhark.CodeGen.RTS.C (backendsCudaH)
    import Futhark.IR.GPUMem hiding
      ( CmpSizeLe,
        GetSize,
        GetSizeMax,
      )
    import Futhark.MonadFreshNames
    import Language.C.Quote.OpenCL qualified as C
    import NeatInterpolation (untrimming)

    mkBoilerplate ::
      T.Text ->
      [(Name, KernelConstExp)] ->
      M.Map Name KernelSafety ->
      [PrimType] ->
      [FailureMsg] ->
      GC.CompilerM OpenCL () ()
    mkBoilerplate :: Text
    -> [(Name, KernelConstExp)]
    -> Map Name KernelSafety
    -> [PrimType]
    -> [FailureMsg]
    -> CompilerM OpenCL () ()
    mkBoilerplate Text
    cuda_program [(Name, KernelConstExp)]
    macros Map Name KernelSafety
    kernels [PrimType]
    types [FailureMsg]
    failures = do
      Text
    -> [(Name, KernelConstExp)]
    -> Text
    -> [Name]
    -> [PrimType]
    -> [FailureMsg]
    -> CompilerM OpenCL () ()
    generateGPUBoilerplate
        Text
    cuda_program
        [(Name, KernelConstExp)]
    macros
        Text
    backendsCudaH
        (Map Name KernelSafety -> [Name]
    forall k a. Map k a -> [k]
    M.keys Map Name KernelSafety
    kernels)
        [PrimType]
    types
        [FailureMsg]
    failures

      HeaderSection -> Definition -> CompilerM OpenCL () ()
    forall op s. HeaderSection -> Definition -> CompilerM op s ()
    GC.headerDecl HeaderSection
    GC.InitDecl [C.cedecl|void futhark_context_config_add_nvrtc_option(struct futhark_context_config *cfg, const char* opt);|]
      HeaderSection -> Definition -> CompilerM OpenCL () ()
    forall op s. HeaderSection -> Definition -> CompilerM op s ()
    GC.headerDecl HeaderSection
    GC.InitDecl [C.cedecl|void futhark_context_config_set_device(struct futhark_context_config *cfg, const char* s);|]
      HeaderSection -> Definition -> CompilerM OpenCL () ()
    forall op s. HeaderSection -> Definition -> CompilerM op s ()
    GC.headerDecl HeaderSection
    GC.InitDecl [C.cedecl|const char* futhark_context_config_get_program(struct futhark_context_config *cfg);|]
      HeaderSection -> Definition -> CompilerM OpenCL () ()
    forall op s. HeaderSection -> Definition -> CompilerM op s ()
    GC.headerDecl HeaderSection
    GC.InitDecl [C.cedecl|void futhark_context_config_set_program(struct futhark_context_config *cfg, const char* s);|]
      HeaderSection -> Definition -> CompilerM OpenCL () ()
    forall op s. HeaderSection -> Definition -> CompilerM op s ()
    GC.headerDecl HeaderSection
    GC.InitDecl [C.cedecl|void futhark_context_config_dump_ptx_to(struct futhark_context_config *cfg, const char* s);|]
      HeaderSection -> Definition -> CompilerM OpenCL () ()
    forall op s. HeaderSection -> Definition -> CompilerM op s ()
    GC.headerDecl HeaderSection
    GC.InitDecl [C.cedecl|void futhark_context_config_load_ptx_from(struct futhark_context_config *cfg, const char* s);|]

    cliOptions :: [Option]
    cliOptions :: [Option]
    cliOptions =
      [Option]
    gpuOptions
        [Option] -> [Option] -> [Option]
    forall a. [a] -> [a] -> [a]
    ++ [ Option
               { optionLongName :: SpaceId
    optionLongName = SpaceId
    "dump-cuda",
                 optionShortName :: Maybe Char
    optionShortName = Maybe Char
    forall a. Maybe a
    Nothing,
                 optionArgument :: OptionArgument
    optionArgument = SpaceId -> OptionArgument
    RequiredArgument SpaceId
    "FILE",
                 optionDescription :: SpaceId
    optionDescription = SpaceId
    "Dump the embedded CUDA kernels to the indicated file.",
                 optionAction :: Stm
    optionAction =
                   [C.cstm|{const char* prog = futhark_context_config_get_program(cfg);
                            if (dump_file(optarg, prog, strlen(prog)) != 0) {
                              fprintf(stderr, "%s: %s\n", optarg, strerror(errno));
                              exit(1);
                            }
                            exit(0);}|]
               },
             Option
               { optionLongName :: SpaceId
    optionLongName = SpaceId
    "load-cuda",
                 optionShortName :: Maybe Char
    optionShortName = Maybe Char
    forall a. Maybe a
    Nothing,
                 optionArgument :: OptionArgument
    optionArgument = SpaceId -> OptionArgument
    RequiredArgument SpaceId
    "FILE",
                 optionDescription :: SpaceId
    optionDescription = SpaceId
    "Instead of using the embedded CUDA kernels, load them from the indicated file.",
                 optionAction :: Stm
    optionAction =
                   [C.cstm|{ size_t n; const char *s = slurp_file(optarg, &n);
                             if (s == NULL) { fprintf(stderr, "%s: %s\n", optarg, strerror(errno)); exit(1); }
                             futhark_context_config_set_program(cfg, s);
                           }|]
               },
             Option
               { optionLongName :: SpaceId
    optionLongName = SpaceId
    "dump-ptx",
                 optionShortName :: Maybe Char
    optionShortName = Maybe Char
    forall a. Maybe a
    Nothing,
                 optionArgument :: OptionArgument
    optionArgument = SpaceId -> OptionArgument
    RequiredArgument SpaceId
    "FILE",
                 optionDescription :: SpaceId
    optionDescription = SpaceId
    "Dump the PTX-compiled version of the embedded kernels to the indicated file.",
                 optionAction :: Stm
    optionAction =
                   [C.cstm|{futhark_context_config_dump_ptx_to(cfg, optarg);
                                         entry_point = NULL;}|]
               },
             Option
               { optionLongName :: SpaceId
    optionLongName = SpaceId
    "load-ptx",
                 optionShortName :: Maybe Char
    optionShortName = Maybe Char
    forall a. Maybe a
    Nothing,
                 optionArgument :: OptionArgument
    optionArgument = SpaceId -> OptionArgument
    RequiredArgument SpaceId
    "FILE",
                 optionDescription :: SpaceId
    optionDescription = SpaceId
    "Load PTX code from the indicated file.",
                 optionAction :: Stm
    optionAction = [C.cstm|futhark_context_config_load_ptx_from(cfg, optarg);|]
               },
             Option
               { optionLongName :: SpaceId
    optionLongName = SpaceId
    "nvrtc-option",
                 optionShortName :: Maybe Char
    optionShortName = Maybe Char
    forall a. Maybe a
    Nothing,
                 optionArgument :: OptionArgument
    optionArgument = SpaceId -> OptionArgument
    RequiredArgument SpaceId
    "OPT",
                 optionDescription :: SpaceId
    optionDescription = SpaceId
    "Add an additional build option to the string passed to NVRTC.",
                 optionAction :: Stm
    optionAction = [C.cstm|futhark_context_config_add_nvrtc_option(cfg, optarg);|]
               }
           ]

    cudaMemoryType :: GC.MemoryType OpenCL ()
    cudaMemoryType :: MemoryType OpenCL ()
    cudaMemoryType SpaceId
    "device" = Type -> CompilerM OpenCL () Type
    forall a. a -> CompilerM OpenCL () a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [C.cty|typename CUdeviceptr|]
    cudaMemoryType SpaceId
    space = MemoryType OpenCL ()
    forall a. HasCallStack => SpaceId -> a
    error MemoryType OpenCL () -> MemoryType OpenCL ()
    forall a b. (a -> b) -> a -> b
    $ SpaceId
    "GPU backend does not support '" SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ SpaceId
    space SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ SpaceId
    "' memory space."

    -- | Compile the program to C with calls to CUDA.
    compileProg :: (MonadFreshNames m) => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts)
    compileProg :: forall (m :: * -> *).
    MonadFreshNames m =>
    Text -> Prog GPUMem -> m (Warnings, CParts)
    compileProg Text
    version Prog GPUMem
    prog = do
      ( ws,
        Program cuda_code cuda_prelude macros kernels types params failures prog'
        ) <-
        Prog GPUMem -> m (Warnings, Program)
    forall (m :: * -> *).
    MonadFreshNames m =>
    Prog GPUMem -> m (Warnings, Program)
    ImpGen.compileProg Prog GPUMem
    prog
      (ws,)
        <$> GC.compileProg
          "cuda"
          version
          params
          operations
          (mkBoilerplate (cuda_prelude <> cuda_code) macros kernels types failures)
          cuda_includes
          (Space "device", [Space "device", DefaultSpace])
          cliOptions
          prog'
      where
        operations :: GC.Operations OpenCL ()
        operations :: Operations OpenCL ()
    operations =
          Operations OpenCL ()
    gpuOperations
            { GC.opsMemoryType = cudaMemoryType,
              GC.opsCritical =
                ( [C.citems|CUDA_SUCCEED_FATAL(cuCtxPushCurrent(ctx->cu_ctx));|],
                  [C.citems|CUDA_SUCCEED_FATAL(cuCtxPopCurrent(&ctx->cu_ctx));|]
                )
            }
        cuda_includes :: Text
    cuda_includes =
          [untrimming|
           #include <cuda.h>
           #include <cuda_runtime.h>
           #include <nvrtc.h>
          |]