-- | Code generation for Python with OpenCL.
module Futhark.CodeGen.Backends.PyOpenCL
  ( compileProg,
  )
where

import Control.Monad
import Control.Monad.Identity
import Data.Map qualified as M
import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericPython hiding (compileProg)
import Futhark.CodeGen.Backends.GenericPython qualified as GP
import Futhark.CodeGen.Backends.GenericPython.AST
import Futhark.CodeGen.Backends.GenericPython.Options
import Futhark.CodeGen.Backends.PyOpenCL.Boilerplate
import Futhark.CodeGen.ImpCode (Count (..))
import Futhark.CodeGen.ImpCode.OpenCL qualified as Imp
import Futhark.CodeGen.ImpGen.OpenCL qualified as ImpGen
import Futhark.CodeGen.RTS.Python (openclPy)
import Futhark.IR.GPUMem (GPUMem, Prog)
import Futhark.MonadFreshNames
import Futhark.Util (zEncodeText)
import Futhark.Util.Pretty (prettyString, prettyText)

-- | Compile the program to Python with calls to OpenCL.
compileProg ::
  (MonadFreshNames m) =>
  CompilerMode ->
  String ->
  Prog GPUMem ->
  m (ImpGen.Warnings, T.Text)
compileProg mode class_name prog = do
  ( ws,
    Imp.Program
      opencl_code
      opencl_prelude
      macros
      kernels
      types
      sizes
      failures
      prog'
    ) <-
    ImpGen.compileProg prog
  -- prepare the strings for assigning the kernels and set them as global
  let assign =
        unlines
          $ map
            ( \x ->
                prettyString $
                  Assign
                    (Var (T.unpack ("self." <> zEncodeText (nameToText x) <> "_var")))
                    (Var $ T.unpack $ "program." <> zEncodeText (nameToText x))
            )
          $ M.keys kernels

  let defines =
        [ Assign (Var "synchronous") $ Bool False,
          Assign (Var "preferred_platform") None,
          Assign (Var "build_options") $ List [],
          Assign (Var "preferred_device") None,
          Assign (Var "default_threshold") None,
          Assign (Var "default_group_size") None,
          Assign (Var "default_num_groups") None,
          Assign (Var "default_tile_size") None,
          Assign (Var "default_reg_tile_size") None,
          Assign (Var "fut_opencl_src") $ RawStringLiteral $ opencl_prelude <> opencl_code
        ]

  let imports =
        [ Import "sys" Nothing,
          Import "numpy" $ Just "np",
          Import "ctypes" $ Just "ct",
          Escape openclPy,
          Import "pyopencl.array" Nothing,
          Import "time" Nothing
        ]

  let constructor =
        Constructor
          [ "self",
            "build_options=build_options",
            "command_queue=None",
            "interactive=False",
            "platform_pref=preferred_platform",
            "device_pref=preferred_device",
            "default_group_size=default_group_size",
            "default_num_groups=default_num_groups",
            "default_tile_size=default_tile_size",
            "default_reg_tile_size=default_reg_tile_size",
            "default_threshold=default_threshold",
            "sizes=sizes"
          ]
          [Escape $ openClInit macros types assign sizes failures]
      options =
        [ Option
            { optionLongName = "platform",
              optionShortName = Just 'p',
              optionArgument = RequiredArgument "str",
              optionAction =
                [Assign (Var "preferred_platform") $ Var "optarg"]
            },
          Option
            { optionLongName = "device",
              optionShortName = Just 'd',
              optionArgument = RequiredArgument "str",
              optionAction =
                [Assign (Var "preferred_device") $ Var "optarg"]
            },
          Option
            { optionLongName = "build-option",
              optionShortName = Nothing,
              optionArgument = RequiredArgument "str",
              optionAction =
                [ Assign (Var "build_options") $
                    BinOp "+" (Var "build_options") $
                      List [Var "optarg"]
                ]
            },
          Option
            { optionLongName = "default-threshold",
              optionShortName = Nothing,
              optionArgument = RequiredArgument "int",
              optionAction =
                [Assign (Var "default_threshold") $ Var "optarg"]
            },
          Option
            { optionLongName = "default-group-size",
              optionShortName = Nothing,
              optionArgument = RequiredArgument "int",
              optionAction =
                [Assign (Var "default_group_size") $ Var "optarg"]
            },
          Option
            { optionLongName = "default-num-groups",
              optionShortName = Nothing,
              optionArgument = RequiredArgument "int",
              optionAction =
                [Assign (Var "default_num_groups") $ Var "optarg"]
            },
          Option
            { optionLongName = "default-tile-size",
              optionShortName = Nothing,
              optionArgument = RequiredArgument "int",
              optionAction =
                [Assign (Var "default_tile_size") $ Var "optarg"]
            },
          Option
            { optionLongName = "default-reg-tile-size",
              optionShortName = Nothing,
              optionArgument = RequiredArgument "int",
              optionAction =
                [Assign (Var "default_reg_tile_size") $ Var "optarg"]
            },
          Option
            { optionLongName = "param",
              optionShortName = Nothing,
              optionArgument = RequiredArgument "param_assignment",
              optionAction =
                [ Assign
                    ( Index
                        (Var "params")
                        ( IdxExp
                            ( Index
                                (Var "optarg")
                                (IdxExp (Integer 0))
                            )
                        )
                    )
                    (Index (Var "optarg") (IdxExp (Integer 1)))
                ]
            }
        ]

  (ws,)
    <$> GP.compileProg
      mode
      class_name
      constructor
      imports
      defines
      operations
      ()
      [Exp $ simpleCall "sync" [Var "self"]]
      options
      prog'
  where
    operations :: Operations Imp.OpenCL ()
    operations =
      Operations
        { opsCompiler = callKernel,
          opsWriteScalar = writeOpenCLScalar,
          opsReadScalar = readOpenCLScalar,
          opsAllocate = allocateOpenCLBuffer,
          opsCopies =
            M.insert (Imp.Space "device", Imp.Space "device") copygpu2gpu $
              opsCopies defaultOperations,
          opsEntryOutput = packArrayOutput,
          opsEntryInput = unpackArrayInput
        }

-- We have many casts to 'long', because PyOpenCL may get confused at
-- the 32-bit numbers that ImpCode uses for offsets and the like.
asLong :: PyExp -> PyExp
asLong x = simpleCall "np.int64" [x]

getParamByKey :: Name -> PyExp
getParamByKey key = Index (Var "self.sizes") (IdxExp $ String $ prettyText key)

kernelConstToExp :: Imp.KernelConst -> PyExp
kernelConstToExp (Imp.SizeConst key _) =
  getParamByKey key
kernelConstToExp (Imp.SizeMaxConst size_class) =
  Var $ "self.max_" <> prettyString size_class

compileConstExp :: Imp.KernelConstExp -> PyExp
compileConstExp e = runIdentity $ compilePrimExp (pure . kernelConstToExp) e

compileBlockDim :: Imp.BlockDim -> CompilerM op s PyExp
compileBlockDim (Left e) = asLong <$> compileExp e
compileBlockDim (Right e) = pure $ compileConstExp e

callKernel :: OpCompiler Imp.OpenCL ()
callKernel (Imp.GetSize v key) = do
  v' <- compileVar v
  stm $ Assign v' $ getParamByKey key
callKernel (Imp.CmpSizeLe v key x) = do
  v' <- compileVar v
  x' <- compileExp x
  stm $ Assign v' $ BinOp "<=" (getParamByKey key) x'
callKernel (Imp.GetSizeMax v size_class) = do
  v' <- compileVar v
  stm $ Assign v' $ kernelConstToExp $ Imp.SizeMaxConst size_class
callKernel (Imp.LaunchKernel safety name shared_memory args num_threadblocks workgroup_size) = do
  num_threadblocks' <- mapM (fmap asLong . compileExp) num_threadblocks
  workgroup_size' <- mapM compileBlockDim workgroup_size
  let kernel_size = zipWith mult_exp num_threadblocks' workgroup_size'
      total_elements = foldl mult_exp (Integer 1) kernel_size
      cond = BinOp "!=" total_elements (Integer 0)
  shared_memory' <- compileExp $ Imp.untyped $ Imp.unCount shared_memory
  body <- collect $ launchKernel name safety kernel_size workgroup_size' shared_memory' args
  stm $ If cond body []

  when (safety >= Imp.SafetyFull) $
    stm $
      Assign (Var "self.failure_is_an_option") $
        compilePrimValue (Imp.IntValue (Imp.Int32Value 1))
  where
    mult_exp = BinOp "*"

launchKernel ::
  Imp.KernelName ->
  Imp.KernelSafety ->
  [PyExp] ->
  [PyExp] ->
  PyExp ->
  [Imp.KernelArg] ->
  CompilerM op s ()
launchKernel kernel_name safety kernel_dims threadblock_dims shared_memory args = do
  let kernel_dims' = Tuple kernel_dims
      threadblock_dims' = Tuple threadblock_dims
      kernel_name' = "self." <> zEncodeText (nameToText kernel_name) <> "_var"
  args' <- mapM processKernelArg args
  let failure_args =
        take
          (Imp.numFailureParams safety)
          [ Var "self.global_failure",
            Var "self.failure_is_an_option",
            Var "self.global_failure_args"
          ]
  stm . Exp $
    simpleCall (T.unpack $ kernel_name' <> ".set_args") $
      [simpleCall "cl.LocalMemory" [simpleCall "max" [shared_memory, Integer 1]]]
        ++ failure_args
        ++ args'
  stm . Exp $
    simpleCall
      "cl.enqueue_nd_range_kernel"
      [Var "self.queue", Var (T.unpack kernel_name'), kernel_dims', threadblock_dims']
  finishIfSynchronous
  where
    processKernelArg :: Imp.KernelArg -> CompilerM op s PyExp
    processKernelArg (Imp.ValueKArg e bt) = toStorage bt <$> compileExp e
    processKernelArg (Imp.MemKArg v) = compileVar v

writeOpenCLScalar :: WriteScalar Imp.OpenCL ()
writeOpenCLScalar mem i bt "device" val = do
  let nparr =
        Call
          (Var "np.array")
          [Arg val, ArgKeyword "dtype" $ Var $ compilePrimType bt]
  stm $
    Exp $
      Call
        (Var "cl.enqueue_copy")
        [ Arg $ Var "self.queue",
          Arg mem,
          Arg nparr,
          ArgKeyword "dst_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt),
          ArgKeyword "is_blocking" $ Var "synchronous"
        ]
writeOpenCLScalar _ _ _ space _ =
  error $ "Cannot write to '" ++ space ++ "' memory space."

readOpenCLScalar :: ReadScalar Imp.OpenCL ()
readOpenCLScalar mem i bt "device" = do
  val <- newVName "read_res"
  let val' = Var $ prettyString val
  let nparr =
        Call
          (Var "np.empty")
          [ Arg $ Integer 1,
            ArgKeyword "dtype" (Var $ compilePrimType bt)
          ]
  stm $ Assign val' nparr
  stm $
    Exp $
      Call
        (Var "cl.enqueue_copy")
        [ Arg $ Var "self.queue",
          Arg val',
          Arg mem,
          ArgKeyword "src_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt),
          ArgKeyword "is_blocking" $ Var "synchronous"
        ]
  stm $ Exp $ simpleCall "sync" [Var "self"]
  pure $ Index val' $ IdxExp $ Integer 0
readOpenCLScalar _ _ _ space =
  error $ "Cannot read from '" ++ space ++ "' memory space."

allocateOpenCLBuffer :: Allocate Imp.OpenCL ()
allocateOpenCLBuffer mem size "device" =
  stm $
    Assign mem $
      simpleCall "opencl_alloc" [Var "self", size, String $ prettyText mem]
allocateOpenCLBuffer _ _ space =
  error $ "Cannot allocate in '" ++ space ++ "' space"

packArrayOutput :: EntryOutput Imp.OpenCL ()
packArrayOutput mem "device" bt ept dims = do
  mem' <- compileVar mem
  dims' <- mapM compileDim dims
  pure $
    Call
      (Var "cl.array.Array")
      [ Arg $ Var "self.queue",
        Arg $ Tuple $ dims' <> [Integer 0 | bt == Imp.Unit],
        Arg $ Var $ compilePrimToExtNp bt ept,
        ArgKeyword "data" mem'
      ]
packArrayOutput _ sid _ _ _ =
  error $ "Cannot return array from " ++ sid ++ " space."

unpackArrayInput :: EntryInput Imp.OpenCL ()
unpackArrayInput mem "device" t s dims e = do
  let type_is_ok =
        BinOp
          "and"
          (BinOp "in" (simpleCall "type" [e]) (List [Var "np.ndarray", Var "cl.array.Array"]))
          (BinOp "==" (Field e "dtype") (Var (compilePrimToExtNp t s)))
  stm $ Assert type_is_ok $ String "Parameter has unexpected type"

  zipWithM_ (unpackDim e) dims [0 ..]

  let memsize' = simpleCall "np.int64" [Field e "nbytes"]
      pyOpenCLArrayCase =
        [Assign mem $ Field e "data"]
  numpyArrayCase <- collect $ do
    allocateOpenCLBuffer mem memsize' "device"
    stm $
      ifNotZeroSize memsize' $
        Exp $
          Call
            (Var "cl.enqueue_copy")
            [ Arg $ Var "self.queue",
              Arg mem,
              Arg $ Call (Var "normaliseArray") [Arg e],
              ArgKeyword "is_blocking" $ Var "synchronous"
            ]

  stm $
    If
      (BinOp "==" (simpleCall "type" [e]) (Var "cl.array.Array"))
      pyOpenCLArrayCase
      numpyArrayCase
unpackArrayInput _ sid _ _ _ _ =
  error $ "Cannot accept array from " ++ sid ++ " space."

ifNotZeroSize :: PyExp -> PyStmt -> PyStmt
ifNotZeroSize e s =
  If (BinOp "!=" e (Integer 0)) [s] []

finishIfSynchronous :: CompilerM op s ()
finishIfSynchronous =
  stm $ If (Var "synchronous") [Exp $ simpleCall "sync" [Var "self"]] []

copygpu2gpu :: DoCopy op s
copygpu2gpu t shape dst (dstoffset, dststride) src (srcoffset, srcstride) = do
  stm . Exp . simpleCall "lmad_copy_gpu2gpu" $
    [ Var "self",
      Var (compilePrimType t),
      dst,
      unCount dstoffset,
      List (map unCount dststride),
      src,
      unCount srcoffset,
      List (map unCount srcstride),
      List (map unCount shape)
    ]