{-# LANGUAGE QuasiQuotes #-}

-- | Code generation for C with OpenCL.
module Futhark.CodeGen.Backends.COpenCL
  ( compileProg,
    GC.CParts (..),
    GC.asLibrary,
    GC.asExecutable,
    GC.asServer,
  )
where

import Control.Monad.State
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.OpenCL qualified as ImpGen
import Futhark.CodeGen.OpenCL.Heuristics
import Futhark.CodeGen.RTS.C (backendsOpenclH)
import Futhark.IR.GPUMem hiding
  ( CmpSizeLe,
    GetSize,
    GetSizeMax,
  )
import Futhark.MonadFreshNames
import Language.C.Quote.OpenCL qualified as C
import Language.C.Syntax qualified as C
import NeatInterpolation (untrimming)

sizeHeuristicsCode :: SizeHeuristic -> C.Stm
sizeHeuristicsCode :: SizeHeuristic -> Stm
sizeHeuristicsCode (SizeHeuristic SpaceId
platform_name DeviceType
device_type WhichSize
which (TPrimExp PrimExp DeviceInfo
what)) =
  [C.cstm|
   if ($exp:which' == 0 &&
       strstr(option->platform_name, $string:platform_name) != NULL &&
       (option->device_type & $exp:(clDeviceType device_type)) == $exp:(clDeviceType device_type)) {
     $items:get_size
   }|]
  where
    clDeviceType :: DeviceType -> Exp
clDeviceType DeviceType
DeviceGPU = [C.cexp|CL_DEVICE_TYPE_GPU|]
    clDeviceType DeviceType
DeviceCPU = [C.cexp|CL_DEVICE_TYPE_CPU|]

    which' :: Exp
which' = case WhichSize
which of
      WhichSize
LockstepWidth -> [C.cexp|ctx->lockstep_width|]
      WhichSize
NumBlocks -> [C.cexp|ctx->cfg->default_num_groups|]
      WhichSize
BlockSize -> [C.cexp|ctx->cfg->default_group_size|]
      WhichSize
TileSize -> [C.cexp|ctx->cfg->default_tile_size|]
      WhichSize
RegTileSize -> [C.cexp|ctx->cfg->default_reg_tile_size|]
      WhichSize
Threshold -> [C.cexp|ctx->cfg->default_threshold|]

    get_size :: [BlockItem]
get_size =
      let (Exp
e, Map SpaceId [BlockItem]
m) = State (Map SpaceId [BlockItem]) Exp
-> Map SpaceId [BlockItem] -> (Exp, Map SpaceId [BlockItem])
forall s a. State s a -> s -> (a, s)
runState ((DeviceInfo -> State (Map SpaceId [BlockItem]) Exp)
-> PrimExp DeviceInfo -> State (Map SpaceId [BlockItem]) Exp
forall (m :: * -> *) v.
Monad m =>
(v -> m Exp) -> PrimExp v -> m Exp
GC.compilePrimExp DeviceInfo -> State (Map SpaceId [BlockItem]) Exp
forall {m :: * -> *}.
MonadState (Map SpaceId [BlockItem]) m =>
DeviceInfo -> m Exp
onLeaf PrimExp DeviceInfo
what) Map SpaceId [BlockItem]
forall a. Monoid a => a
mempty
       in [[BlockItem]] -> [BlockItem]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Map SpaceId [BlockItem] -> [[BlockItem]]
forall k a. Map k a -> [a]
M.elems Map SpaceId [BlockItem]
m) [BlockItem] -> [BlockItem] -> [BlockItem]
forall a. [a] -> [a] -> [a]
++ [[C.citem|$exp:which' = $exp:e;|]]

    onLeaf :: DeviceInfo -> m Exp
onLeaf (DeviceInfo SpaceId
s) = do
      let s' :: SpaceId
s' = SpaceId
"CL_DEVICE_" SpaceId -> SpaceId -> SpaceId
forall a. [a] -> [a] -> [a]
++ SpaceId
s
          v :: SpaceId
v = SpaceId
s SpaceId -> SpaceId -> SpaceId
forall a. [a] -> [a] -> [a]
++ SpaceId
"_val"
      Map SpaceId [BlockItem]
m <- m (Map SpaceId [BlockItem])
forall s (m :: * -> *). MonadState s m => m s
get
      case SpaceId -> Map SpaceId [BlockItem] -> Maybe [BlockItem]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup SpaceId
s Map SpaceId [BlockItem]
m of
        Maybe [BlockItem]
Nothing ->
          -- XXX: Cheating with the type here; works for the infos we
          -- currently use because we zero-initialise and assume a
          -- little-endian platform, but should be made more
          -- size-aware in the future.
          (Map SpaceId [BlockItem] -> Map SpaceId [BlockItem]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map SpaceId [BlockItem] -> Map SpaceId [BlockItem]) -> m ())
-> (Map SpaceId [BlockItem] -> Map SpaceId [BlockItem]) -> m ()
forall a b. (a -> b) -> a -> b
$
            SpaceId
-> [BlockItem]
-> Map SpaceId [BlockItem]
-> Map SpaceId [BlockItem]
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert
              SpaceId
s'
              [C.citems|size_t $id:v = 0;
                        clGetDeviceInfo(ctx->device, $id:s',
                                        sizeof($id:v), &$id:v,
                                        NULL);|]
        Just [BlockItem]
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

      Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cexp|$id:v|]

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
opencl_program [(Name, KernelConstExp)]
macros Map Name KernelSafety
kernels [PrimType]
types [FailureMsg]
failures = do
  Text
-> [(Name, KernelConstExp)]
-> Text
-> [Name]
-> [PrimType]
-> [FailureMsg]
-> CompilerM OpenCL () ()
generateGPUBoilerplate
    Text
opencl_program
    [(Name, KernelConstExp)]
macros
    Text
backendsOpenclH
    (Map Name KernelSafety -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name KernelSafety
kernels)
    [PrimType]
types
    [FailureMsg]
failures

  Definition -> CompilerM OpenCL () ()
forall op s. Definition -> CompilerM op s ()
GC.earlyDecl
    [C.cedecl|void post_opencl_setup(struct futhark_context *ctx, struct opencl_device_option *option) {
             $stms:(map sizeHeuristicsCode sizeHeuristicsTable)
             }|]

  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|void futhark_context_config_set_platform(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_select_device_interactively(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_list_devices(struct futhark_context_config *cfg);|]
  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_binary_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_binary_from(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_set_command_queue(struct futhark_context_config *cfg, typename cl_command_queue);|]
  HeaderSection -> Definition -> CompilerM OpenCL () ()
forall op s. HeaderSection -> Definition -> CompilerM op s ()
GC.headerDecl HeaderSection
GC.MiscDecl [C.cedecl|typename cl_command_queue futhark_context_get_command_queue(struct futhark_context* ctx);|]

cliOptions :: [Option]
cliOptions :: [Option]
cliOptions =
  [Option]
gpuOptions
    [Option] -> [Option] -> [Option]
forall a. [a] -> [a] -> [a]
++ [ Option
           { optionLongName :: SpaceId
optionLongName = SpaceId
"platform",
             optionShortName :: Maybe Char
optionShortName = Char -> Maybe Char
forall a. a -> Maybe a
Just Char
'p',
             optionArgument :: OptionArgument
optionArgument = SpaceId -> OptionArgument
RequiredArgument SpaceId
"NAME",
             optionDescription :: SpaceId
optionDescription = SpaceId
"Use the first OpenCL platform whose name contains the given string.",
             optionAction :: Stm
optionAction = [C.cstm|futhark_context_config_set_platform(cfg, optarg);|]
           },
         Option
           { optionLongName :: SpaceId
optionLongName = SpaceId
"dump-opencl",
             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 OpenCL program 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-opencl",
             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 OpenCL program, load it 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-opencl-binary",
             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 compiled version of the embedded OpenCL program to the indicated file.",
             optionAction :: Stm
optionAction =
               [C.cstm|{futhark_context_config_dump_binary_to(cfg, optarg);
                                     entry_point = NULL;}|]
           },
         Option
           { optionLongName :: SpaceId
optionLongName = SpaceId
"load-opencl-binary",
             optionShortName :: Maybe Char
optionShortName = Maybe Char
forall a. Maybe a
Nothing,
             optionArgument :: OptionArgument
optionArgument = SpaceId -> OptionArgument
RequiredArgument SpaceId
"FILE",
             optionDescription :: SpaceId
optionDescription = SpaceId
"Load an OpenCL binary from the indicated file.",
             optionAction :: Stm
optionAction = [C.cstm|futhark_context_config_load_binary_from(cfg, optarg);|]
           },
         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 clBuildProgram().",
             optionAction :: Stm
optionAction = [C.cstm|futhark_context_config_add_build_option(cfg, optarg);|]
           },
         Option
           { optionLongName :: SpaceId
optionLongName = SpaceId
"list-devices",
             optionShortName :: Maybe Char
optionShortName = Maybe Char
forall a. Maybe a
Nothing,
             optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
             optionDescription :: SpaceId
optionDescription = SpaceId
"List all OpenCL devices and platforms available on the system.",
             optionAction :: Stm
optionAction =
               [C.cstm|{futhark_context_config_list_devices(cfg);
                        entry_point = NULL;}|]
           }
       ]

openclMemoryType :: GC.MemoryType OpenCL ()
openclMemoryType :: MemoryType OpenCL ()
openclMemoryType SpaceId
"device" = Type -> CompilerM OpenCL () Type
forall a. a -> CompilerM OpenCL () a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cty|typename cl_mem|]
openclMemoryType 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 OpenCL.
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
opencl_code Text
opencl_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
"opencl"
      Text
version
      ParamMap
params
      Operations OpenCL ()
operations
      (Text
-> [(Name, KernelConstExp)]
-> Map Name KernelSafety
-> [PrimType]
-> [FailureMsg]
-> CompilerM OpenCL () ()
mkBoilerplate (Text
opencl_prelude Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
opencl_code) [(Name, KernelConstExp)]
macros Map Name KernelSafety
kernels [PrimType]
types [FailureMsg]
failures)
      Text
opencl_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
        { opsMemoryType :: MemoryType OpenCL ()
GC.opsMemoryType = MemoryType OpenCL ()
openclMemoryType
        }
    opencl_includes :: Text
opencl_includes =
      [untrimming|
       #define CL_TARGET_OPENCL_VERSION 120
       #define CL_USE_DEPRECATED_OPENCL_1_2_APIS
       #ifdef __APPLE__
       #define CL_SILENCE_DEPRECATION
       #include <OpenCL/cl.h>
       #else
       #include <CL/cl.h>
       #endif
       |]