{-# LANGUAGE QuasiQuotes #-}

-- | Code generation for HIP.
module Futhark.CodeGen.Backends.HIP
  ( 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.HIP qualified as ImpGen
import Futhark.CodeGen.RTS.C (backendsHipH)
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
hip_program [(Name, KernelConstExp)]
macros Map Name KernelSafety
kernels [PrimType]
types [FailureMsg]
failures = do
  Text
-> [(Name, KernelConstExp)]
-> Text
-> [Name]
-> [PrimType]
-> [FailureMsg]
-> CompilerM OpenCL () ()
generateGPUBoilerplate
    Text
hip_program
    [(Name, KernelConstExp)]
macros
    Text
backendsHipH
    (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_build_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);|]

cliOptions :: [Option]
cliOptions :: [Option]
cliOptions =
  [Option]
gpuOptions
    [Option] -> [Option] -> [Option]
forall a. [a] -> [a] -> [a]
++ [ Option
           { optionLongName :: SpaceId
optionLongName = SpaceId
"dump-hip",
             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 HIP 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-hip",
             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 HIP 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
"build-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 HIPRTC.",
             optionAction :: Stm
optionAction = [C.cstm|futhark_context_config_add_build_option(cfg, optarg);|]
           }
       ]

hipMemoryType :: GC.MemoryType OpenCL ()
hipMemoryType :: MemoryType OpenCL ()
hipMemoryType SpaceId
"device" = Type -> CompilerM OpenCL () Type
forall a. a -> CompilerM OpenCL () a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cty|typename hipDeviceptr_t|]
hipMemoryType 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 HIP.
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
  ( Warnings
ws,
    Program Text
hip_code Text
hip_prelude [(Name, KernelConstExp)]
macros Map Name KernelSafety
kernels [PrimType]
types ParamMap
params [FailureMsg]
failures Definitions OpenCL
prog'
    ) <-
    Prog GPUMem -> m (Warnings, Program)
forall (m :: * -> *).
MonadFreshNames m =>
Prog GPUMem -> m (Warnings, Program)
ImpGen.compileProg Prog GPUMem
prog
  (Warnings
ws,)
    (CParts -> (Warnings, CParts)) -> m CParts -> m (Warnings, CParts)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text
-> Text
-> ParamMap
-> Operations OpenCL ()
-> CompilerM OpenCL () ()
-> Text
-> (Space, [Space])
-> [Option]
-> Definitions OpenCL
-> m CParts
forall (m :: * -> *) op.
MonadFreshNames m =>
Text
-> Text
-> ParamMap
-> Operations op ()
-> CompilerM op () ()
-> Text
-> (Space, [Space])
-> [Option]
-> Definitions op
-> m CParts
GC.compileProg
      Text
"hip"
      Text
version
      ParamMap
params
      Operations OpenCL ()
operations
      (Text
-> [(Name, KernelConstExp)]
-> Map Name KernelSafety
-> [PrimType]
-> [FailureMsg]
-> CompilerM OpenCL () ()
mkBoilerplate (Text
hip_prelude Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
hip_code) [(Name, KernelConstExp)]
macros Map Name KernelSafety
kernels [PrimType]
types [FailureMsg]
failures)
      Text
hip_includes
      (SpaceId -> Space
Space SpaceId
"device", [SpaceId -> Space
Space SpaceId
"device", Space
DefaultSpace])
      [Option]
cliOptions
      Definitions OpenCL
prog'
  where
    operations :: GC.Operations OpenCL ()
    operations :: Operations OpenCL ()
operations =
      Operations OpenCL ()
gpuOperations
        { GC.opsMemoryType = hipMemoryType
        }
    hip_includes :: Text
hip_includes =
      [untrimming|
       #define __HIP_PLATFORM_AMD__
       #include <hip/hip_runtime.h>
       #include <hip/hiprtc.h>
      |]