{-# LANGUAGE CPP                 #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE QuasiQuotes         #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.Embed
-- Copyright   : [2017..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.LLVM.Embed (

  Embed(..),
  embedAfun, embedOpenAfun,
  embedOpenAcc,

) where

import LLVM.AST.Type.Name

import Data.Array.Accelerate.AST                                    ( PreOpenAfun(..), ArrayVar, Direction(..), Exp, liftALeftHandSide, liftOpenExp, arrayR )
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Elt
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Type

import Data.Array.Accelerate.LLVM.AST
import Data.Array.Accelerate.LLVM.CodeGen.Environment
import Data.Array.Accelerate.LLVM.Compile
import Data.Array.Accelerate.LLVM.Link

import Data.ByteString.Short                                        ( ShortByteString )
import GHC.Ptr                                                      ( Ptr(..) )
import Language.Haskell.TH                                          ( Q, TExp )
import System.IO.Unsafe
import qualified Data.ByteString.Short.Internal                     as BS
import qualified Language.Haskell.TH                                as TH
import qualified Language.Haskell.TH.Syntax                         as TH

#if MIN_VERSION_containers(0,5,9)
import qualified Data.IntMap.Internal                               as IM
#elif MIN_VERSION_containers(0,5,8)
import qualified Data.IntMap.Base                                   as IM
#else
import qualified Data.IntMap                                        as IM
#endif


class Embed arch where

  -- | Turn the compiled object into a TemplateHaskell expression, suitable for
  -- use in a splice. The splice should evaluate into the backend-specific
  -- executable representation.
  --
  embedForTarget
      :: arch
      -> ObjectR arch
      -> Q (TExp (ExecutableR arch))


-- | Embed the compiled array function into a TemplateHaskell expression,
-- suitable for use in a splice.
--
{-# INLINEABLE embedAfun #-}
embedAfun
    :: Embed arch
    => arch
    -> CompiledAfun arch f
    -> Q (TExp (ExecAfun arch f))
embedAfun :: arch -> CompiledAfun arch f -> Q (TExp (ExecAfun arch f))
embedAfun = arch -> CompiledAfun arch f -> Q (TExp (ExecAfun arch f))
forall arch aenv f.
(HasCallStack, Embed arch) =>
arch
-> CompiledOpenAfun arch aenv f
-> Q (TExp (ExecOpenAfun arch aenv f))
embedOpenAfun

{-# INLINEABLE embedOpenAfun #-}
embedOpenAfun
    :: (HasCallStack, Embed arch)
    => arch
    -> CompiledOpenAfun arch aenv f
    -> Q (TExp (ExecOpenAfun arch aenv f))
embedOpenAfun :: arch
-> CompiledOpenAfun arch aenv f
-> Q (TExp (ExecOpenAfun arch aenv f))
embedOpenAfun arch
arch (Alam ALeftHandSide a aenv aenv'
lhs PreOpenAfun (CompiledOpenAcc arch) aenv' t1
l) = [|| Alam $$(liftALeftHandSide lhs) $$(embedOpenAfun arch l) ||]
embedOpenAfun arch
arch (Abody CompiledOpenAcc arch aenv f
b)    = [|| Abody $$(embedOpenAcc arch b) ||]

{-# INLINEABLE embedOpenAcc #-}
embedOpenAcc
    :: forall arch aenv arrs. (HasCallStack, Embed arch)
    => arch
    -> CompiledOpenAcc arch aenv arrs
    -> Q (TExp (ExecOpenAcc arch aenv arrs))
embedOpenAcc :: arch
-> CompiledOpenAcc arch aenv arrs
-> Q (TExp (ExecOpenAcc arch aenv arrs))
embedOpenAcc arch
arch = CompiledOpenAcc arch aenv arrs
-> Q (TExp (ExecOpenAcc arch aenv arrs))
forall aenv' arrs'.
CompiledOpenAcc arch aenv' arrs'
-> Q (TExp (ExecOpenAcc arch aenv' arrs'))
liftA
  where
    liftA :: CompiledOpenAcc arch aenv' arrs' -> Q (TExp (ExecOpenAcc arch aenv' arrs'))
    liftA :: CompiledOpenAcc arch aenv' arrs'
-> Q (TExp (ExecOpenAcc arch aenv' arrs'))
liftA CompiledOpenAcc arch aenv' arrs'
acc = case CompiledOpenAcc arch aenv' arrs'
acc of
        PlainAcc ArraysR arrs'
repr PreOpenAccCommand CompiledOpenAcc arch aenv' arrs'
pacc          -> [|| EvalAcc $$(liftArraysR repr) $$(liftPreOpenAccCommand arch pacc) ||]
        BuildAcc ArraysR arrs'
repr Gamma aenv'
aenv ObjectR arch
obj PreOpenAccSkeleton CompiledOpenAcc arch aenv' arrs'
pacc -> [|| ExecAcc $$(liftArraysR repr) $$(liftGamma aenv) $$(embedForTarget arch obj) $$(liftPreOpenAccSkeleton arch pacc) ||]

    liftGamma :: Gamma aenv' -> Q (TExp (Gamma aenv'))
#if MIN_VERSION_containers(0,5,8)
    liftGamma :: Gamma aenv' -> Q (TExp (Gamma aenv'))
liftGamma Gamma aenv'
IM.Nil           = [|| IM.Nil ||]
    liftGamma (IM.Bin Prefix
p Prefix
m Gamma aenv'
l Gamma aenv'
r) = [|| IM.Bin p m $$(liftGamma l) $$(liftGamma r) ||]
    liftGamma (IM.Tip Prefix
k (Label, Idx' aenv')
v)     = [|| IM.Tip k $$(liftV v) ||]
#else
    -- O(n) at runtime to reconstruct the set
    liftGamma aenv             = [|| IM.fromAscList $$(liftIM (IM.toAscList aenv)) ||]
      where
        liftIM :: [(Int, (Label, Idx' aenv'))] -> Q (TExp [(Int, (Label, Idx' aenv'))])
        liftIM im =
          TH.TExp . TH.ListE <$> mapM (\(k,v) -> TH.unTypeQ [|| (k, $$(liftV v)) ||]) im
#endif
    liftV :: (Label, Idx' aenv') -> Q (TExp (Label, Idx' aenv'))
    liftV :: (Label, Idx' aenv') -> Q (TExp (Label, Idx' aenv'))
liftV (Label ShortByteString
n, Idx' ArrayR (Array sh e)
repr Idx aenv' (Array sh e)
ix) = [|| (Label $$(liftSBS n), Idx' $$(liftArrayR repr) $$(liftIdx ix)) ||]

    -- O(n) at runtime to copy from the Addr# to the ByteArray#. We should
    -- be able to do this without copying, but I don't think the definition of
    -- ByteArray# is exported (or it is deeply magical).
    liftSBS :: ShortByteString -> Q (TExp ShortByteString)
    liftSBS :: ShortByteString -> Q (TExp ShortByteString)
liftSBS ShortByteString
bs =
      let bytes :: [Word8]
bytes = ShortByteString -> [Word8]
BS.unpack ShortByteString
bs
          len :: Prefix
len   = ShortByteString -> Prefix
BS.length ShortByteString
bs
      in
      [|| unsafePerformIO $ BS.createFromPtr $$( TH.unsafeTExpCoerce [| Ptr $(TH.litE (TH.StringPrimL bytes)) |]) len ||]


{-# INLINEABLE liftPreOpenAfun #-}
liftPreOpenAfun
    :: (HasCallStack, Embed arch)
    => arch
    -> PreOpenAfun (CompiledOpenAcc arch) aenv t
    -> Q (TExp (PreOpenAfun (ExecOpenAcc arch) aenv t))
liftPreOpenAfun :: arch
-> PreOpenAfun (CompiledOpenAcc arch) aenv t
-> Q (TExp (PreOpenAfun (ExecOpenAcc arch) aenv t))
liftPreOpenAfun arch
arch (Alam ALeftHandSide a aenv aenv'
lhs PreOpenAfun (CompiledOpenAcc arch) aenv' t1
f) = [|| Alam $$(liftALeftHandSide lhs) $$(liftPreOpenAfun arch f) ||]
liftPreOpenAfun arch
arch (Abody CompiledOpenAcc arch aenv t
b)    = [|| Abody $$(embedOpenAcc arch b) ||]

{-# INLINEABLE liftPreOpenAccCommand #-}
liftPreOpenAccCommand
    :: forall arch aenv a. (HasCallStack, Embed arch)
    => arch
    -> PreOpenAccCommand CompiledOpenAcc arch aenv a
    -> Q (TExp (PreOpenAccCommand ExecOpenAcc arch aenv a))
liftPreOpenAccCommand :: arch
-> PreOpenAccCommand CompiledOpenAcc arch aenv a
-> Q (TExp (PreOpenAccCommand ExecOpenAcc arch aenv a))
liftPreOpenAccCommand arch
arch PreOpenAccCommand CompiledOpenAcc arch aenv a
pacc =
  let
      liftA :: CompiledOpenAcc arch aenv' arrs -> Q (TExp (ExecOpenAcc arch aenv' arrs))
      liftA :: CompiledOpenAcc arch aenv' arrs
-> Q (TExp (ExecOpenAcc arch aenv' arrs))
liftA = arch
-> CompiledOpenAcc arch aenv' arrs
-> Q (TExp (ExecOpenAcc arch aenv' arrs))
forall arch aenv arrs.
(HasCallStack, Embed arch) =>
arch
-> CompiledOpenAcc arch aenv arrs
-> Q (TExp (ExecOpenAcc arch aenv arrs))
embedOpenAcc arch
arch

      liftE :: Exp aenv t -> Q (TExp (Exp aenv t))
      liftE :: Exp aenv t -> Q (TExp (Exp aenv t))
liftE = Exp aenv t -> Q (TExp (Exp aenv t))
forall env aenv t.
OpenExp env aenv t -> Q (TExp (OpenExp env aenv t))
liftOpenExp

      liftAF :: PreOpenAfun (CompiledOpenAcc arch) aenv f -> Q (TExp (PreOpenAfun (ExecOpenAcc arch) aenv f))
      liftAF :: PreOpenAfun (CompiledOpenAcc arch) aenv f
-> Q (TExp (PreOpenAfun (ExecOpenAcc arch) aenv f))
liftAF = arch
-> PreOpenAfun (CompiledOpenAcc arch) aenv f
-> Q (TExp (PreOpenAfun (ExecOpenAcc arch) aenv f))
forall arch aenv f.
(HasCallStack, Embed arch) =>
arch
-> CompiledOpenAfun arch aenv f
-> Q (TExp (ExecOpenAfun arch aenv f))
liftPreOpenAfun arch
arch
  in
  case PreOpenAccCommand CompiledOpenAcc arch aenv a
pacc of
    Avar ArrayVar aenv a
v            -> [|| Avar $$(liftArrayVar v) ||]
    Alet ALeftHandSide bnd aenv aenv'
lhs CompiledOpenAcc arch aenv bnd
bnd CompiledOpenAcc arch aenv' a
body -> [|| Alet $$(liftALeftHandSide lhs) $$(liftA bnd) $$(liftA body) ||]
    Alloc ArrayR (Array sh e)
repr Exp aenv sh
sh     -> [|| Alloc $$(liftArrayR repr) $$(liftE sh) ||]
    Use ArrayR (Array sh e)
repr Array sh e
a        -> [|| Use $$(liftArrayR repr) $$(liftArray repr a) ||]
    Unit TypeR e
tp Exp aenv e
e         -> [|| Unit $$(liftTypeR tp) $$(liftE e) ||]
    Apair CompiledOpenAcc arch aenv arrs1
a1 CompiledOpenAcc arch aenv arrs2
a2       -> [|| Apair $$(liftA a1) $$(liftA a2) ||]
    PreOpenAccCommand CompiledOpenAcc arch aenv a
Anil              -> [|| Anil ||]
    Apply ArraysR a
repr PreOpenAfun (CompiledOpenAcc arch) aenv (as -> a)
f CompiledOpenAcc arch aenv as
a    -> [|| Apply $$(liftArraysR repr) $$(liftAF f) $$(liftA a) ||]
    Acond Exp aenv Word8
p CompiledOpenAcc arch aenv a
t CompiledOpenAcc arch aenv a
e       -> [|| Acond $$(liftE p) $$(liftA t) $$(liftA e) ||]
    Awhile PreOpenAfun (CompiledOpenAcc arch) aenv (a -> Scalar Word8)
p PreOpenAfun (CompiledOpenAcc arch) aenv (a -> a)
f CompiledOpenAcc arch aenv a
a      -> [|| Awhile $$(liftAF p) $$(liftAF f) $$(liftA a) ||]
    Reshape ShapeR sh
shr Exp aenv sh
sh ArrayVar aenv (Array sh' e)
v  -> [|| Reshape $$(liftShapeR shr) $$(liftE sh) $$(liftArrayVar v) ||]
    Unzip UnzipIdx tup e
tix ArrayVar aenv (Array sh tup)
v       -> [|| Unzip $$(liftUnzipIdx tix) $$(liftArrayVar v) ||]
    Aforeign{}        -> String -> Q (TExp (PreOpenAccCommand ExecOpenAcc arch aenv a))
forall a. HasCallStack => String -> a
internalError String
"using foreign functions from template-haskell is not supported yet"

{-# INLINEABLE liftPreOpenAccSkeleton #-}
liftPreOpenAccSkeleton
    :: forall arch aenv a. (HasCallStack, Embed arch)
    => arch
    -> PreOpenAccSkeleton CompiledOpenAcc arch aenv a
    -> Q (TExp (PreOpenAccSkeleton ExecOpenAcc arch aenv a))
liftPreOpenAccSkeleton :: arch
-> PreOpenAccSkeleton CompiledOpenAcc arch aenv a
-> Q (TExp (PreOpenAccSkeleton ExecOpenAcc arch aenv a))
liftPreOpenAccSkeleton arch
arch PreOpenAccSkeleton CompiledOpenAcc arch aenv a
pacc =
  let
      liftA :: CompiledOpenAcc arch aenv arrs -> Q (TExp (ExecOpenAcc arch aenv arrs))
      liftA :: CompiledOpenAcc arch aenv arrs
-> Q (TExp (ExecOpenAcc arch aenv arrs))
liftA = arch
-> CompiledOpenAcc arch aenv arrs
-> Q (TExp (ExecOpenAcc arch aenv arrs))
forall arch aenv arrs.
(HasCallStack, Embed arch) =>
arch
-> CompiledOpenAcc arch aenv arrs
-> Q (TExp (ExecOpenAcc arch aenv arrs))
embedOpenAcc arch
arch

      liftD :: DelayedOpenAcc CompiledOpenAcc arch aenv arrs -> Q (TExp (DelayedOpenAcc ExecOpenAcc arch aenv arrs))
      liftD :: DelayedOpenAcc CompiledOpenAcc arch aenv arrs
-> Q (TExp (DelayedOpenAcc ExecOpenAcc arch aenv arrs))
liftD (Delayed ArrayR (Array sh e)
repr Exp aenv sh
sh) = [|| Delayed $$(liftArrayR repr) $$(liftE sh) ||]
      liftD (Manifest ArraysR (Array sh e)
repr CompiledOpenAcc arch aenv (Array sh e)
a) = [|| Manifest $$(liftArraysR repr) $$(liftA a) ||]

      liftE :: Exp aenv t -> Q (TExp (Exp aenv t))
      liftE :: Exp aenv t -> Q (TExp (Exp aenv t))
liftE = Exp aenv t -> Q (TExp (Exp aenv t))
forall env aenv t.
OpenExp env aenv t -> Q (TExp (OpenExp env aenv t))
liftOpenExp

      liftS :: ShapeR sh -> sh -> Q (TExp sh)
      liftS :: ShapeR sh -> sh -> Q (TExp sh)
liftS ShapeR sh
shr sh
sh = [|| $$(liftElt (shapeType shr) sh) ||]

      liftZ :: HasInitialValue -> Q (TExp HasInitialValue)
      liftZ :: HasInitialValue -> Q (TExp HasInitialValue)
liftZ HasInitialValue
True  = [|| True  ||]
      liftZ HasInitialValue
False = [|| False ||]

      liftDir :: Direction -> Q (TExp Direction)
      liftDir :: Direction -> Q (TExp Direction)
liftDir Direction
LeftToRight = [|| LeftToRight ||]
      liftDir Direction
RightToLeft = [|| RightToLeft ||]
  in
  case PreOpenAccSkeleton CompiledOpenAcc arch aenv a
pacc of
    Map TypeR b
tp CompiledOpenAcc arch aenv (Array sh a)
a             -> [|| Map $$(liftTypeR tp) $$(liftA a) ||]
    Generate ArrayR (Array sh e)
repr Exp aenv sh
sh     -> [|| Generate $$(liftArrayR repr) $$(liftE sh) ||]
    Transform ArrayR (Array sh' b)
repr Exp aenv sh'
sh CompiledOpenAcc arch aenv (Array sh a)
a  -> [|| Transform $$(liftArrayR repr) $$(liftE sh) $$(liftA a) ||]
    Backpermute ShapeR sh'
shr Exp aenv sh'
sh CompiledOpenAcc arch aenv (Array sh e)
a -> [|| Backpermute $$(liftShapeR shr) $$(liftE sh) $$(liftA a) ||]
    Fold HasInitialValue
z DelayedOpenAcc CompiledOpenAcc arch aenv (Array (sh, Prefix) e)
a             -> [|| Fold $$(liftZ z) $$(liftD a) ||]
    FoldSeg IntegralType i
i HasInitialValue
z DelayedOpenAcc CompiledOpenAcc arch aenv (Array (sh, Prefix) e)
a DelayedOpenAcc CompiledOpenAcc arch aenv (Segments i)
s      -> [|| FoldSeg $$(liftIntegralType i) $$(liftZ z) $$(liftD a) $$(liftD s) ||]
    Scan Direction
d HasInitialValue
z DelayedOpenAcc CompiledOpenAcc arch aenv (Array (sh, Prefix) e)
a           -> [|| Scan $$(liftDir d) $$(liftZ z) $$(liftD a) ||]
    Scan' Direction
d DelayedOpenAcc CompiledOpenAcc arch aenv (Array (sh, Prefix) e)
a            -> [|| Scan' $$(liftDir d) $$(liftD a) ||]
    Permute CompiledOpenAcc arch aenv (Array sh' e)
d DelayedOpenAcc CompiledOpenAcc arch aenv (Array sh e)
a          -> [|| Permute $$(liftA d) $$(liftD a) ||]
    Stencil1 TypeR b
tp sh
h DelayedOpenAcc CompiledOpenAcc arch aenv (Array sh a)
a      -> [|| Stencil1 $$(liftTypeR tp) $$(liftS (arrayRshape $ arrayR a) h) $$(liftD a) ||]
    Stencil2 TypeR c
tp sh
h DelayedOpenAcc CompiledOpenAcc arch aenv (Array sh a)
a DelayedOpenAcc CompiledOpenAcc arch aenv (Array sh b)
b    -> [|| Stencil2 $$(liftTypeR tp) $$(liftS (arrayRshape $ arrayR a) h) $$(liftD a) $$(liftD b) ||]

liftArrayVar :: ArrayVar aenv v -> Q (TExp (ArrayVar aenv v))
liftArrayVar :: ArrayVar aenv v -> Q (TExp (ArrayVar aenv v))
liftArrayVar (Var ArrayR v
tp Idx aenv v
v) = [|| Var $$(liftArrayR tp) $$(liftIdx v) ||]

liftUnzipIdx :: UnzipIdx tup e -> Q (TExp (UnzipIdx tup e))
liftUnzipIdx :: UnzipIdx tup e -> Q (TExp (UnzipIdx tup e))
liftUnzipIdx UnzipIdx tup e
UnzipId                    = [|| UnzipId ||]
liftUnzipIdx (UnzipPrj PairIdx tup b
PairIdxLeft  UnzipIdx b e
ix) = [|| UnzipPrj PairIdxLeft  $$(liftUnzipIdx ix) ||]
liftUnzipIdx (UnzipPrj PairIdx tup b
PairIdxRight UnzipIdx b e
ix) = [|| UnzipPrj PairIdxRight $$(liftUnzipIdx ix) ||]
liftUnzipIdx UnzipIdx tup e
UnzipUnit                  = [|| UnzipUnit ||]
liftUnzipIdx (UnzipPair UnzipIdx tup b1
ix1 UnzipIdx tup b2
ix2)        = [|| UnzipPair $$(liftUnzipIdx ix1) $$(liftUnzipIdx ix2) ||]