content: {-# LANGUAGE EmptyCase           #-}
  {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE OverloadedStrings   #-}
  {-# LANGUAGE RecordWildCards     #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections       #-}
  {-# LANGUAGE TypeApplications    #-}
  {-# LANGUAGE TypeFamilies        #-}
  {-# LANGUAGE TypeOperators       #-}
  {-# OPTIONS_HADDOCK hide #-}
  -- |
  -- Module      : Data.Array.Accelerate.LLVM.Compile
  -- Copyright   : [2014..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.Compile (
  
    Compile(..),
    compileAcc, compileAfun,
  
    CompiledOpenAcc(..), CompiledOpenAfun,
    CompiledAcc, CompiledAfun,
  
  ) where
  
  import Data.Array.Accelerate.AST
  import Data.Array.Accelerate.AST.Environment
  import Data.Array.Accelerate.AST.Idx
  import Data.Array.Accelerate.AST.LeftHandSide
  import Data.Array.Accelerate.AST.Var
  import Data.Array.Accelerate.Analysis.Match
  import Data.Array.Accelerate.Error
  import Data.Array.Accelerate.Representation.Array
  import Data.Array.Accelerate.Representation.Shape
  import Data.Array.Accelerate.Representation.Stencil
  import Data.Array.Accelerate.Representation.Tag
  import Data.Array.Accelerate.Representation.Type
  import Data.Array.Accelerate.Trafo.Delayed
  import Data.Array.Accelerate.Trafo.Substitution
  import Data.Array.Accelerate.Type
  import qualified Data.Array.Accelerate.Sugar.Foreign                as A
  
  import Data.Array.Accelerate.LLVM.CodeGen.Environment
  import Data.Array.Accelerate.LLVM.Foreign
  import Data.Array.Accelerate.LLVM.State
  import qualified Data.Array.Accelerate.LLVM.AST                     as AST
  
  import Data.IntMap                                                  ( IntMap )
  import Formatting                                                   ( (%) )
  import Control.Applicative                                          hiding ( Const )
  import Prelude                                                      hiding ( map, unzip, zipWith, scanl, scanl1, scanr, scanr1, exp )
  import qualified Data.IntMap                                        as IntMap
  
  
  class Foreign arch => Compile arch where
    data ObjectR arch
  
    -- | Compile an accelerate computation into some backend-specific code that
    -- will be used to execute the given array expression. The code is not yet
    -- linked into the running executable.
    --
    compileForTarget
        :: PreOpenAcc DelayedOpenAcc aenv a
        -> Gamma aenv
        -> LLVM arch (ObjectR arch)
  
  
  data CompiledOpenAcc arch aenv a where
    BuildAcc  :: ArraysR a
              -> Gamma aenv
              -> ObjectR arch
              -> AST.PreOpenAccSkeleton CompiledOpenAcc arch aenv a
              -> CompiledOpenAcc arch aenv a
  
    PlainAcc  :: ArraysR a
              -> AST.PreOpenAccCommand  CompiledOpenAcc arch aenv a
              -> CompiledOpenAcc arch aenv a
  
  
  -- An annotated AST with embedded build products
  --
  type CompiledOpenAfun arch  = PreOpenAfun (CompiledOpenAcc arch)
  
  type CompiledAcc arch a     = CompiledOpenAcc arch () a
  type CompiledAfun arch a    = CompiledOpenAfun arch () a
  
  
  -- | Generate and compile code for an array expression. The returned expression
  -- is annotated with the compilation products required to executed each
  -- operation on the given target, together with the list of array variables
  -- referenced from the embedded scalar expressions.
  --
  {-# INLINEABLE compileAcc #-}
  compileAcc
      :: (HasCallStack, Compile arch)
      => DelayedAcc a
      -> LLVM arch (CompiledAcc arch a)
  compileAcc = compileOpenAcc
  
  {-# INLINEABLE compileAfun #-}
  compileAfun
      :: (HasCallStack, Compile arch)
      => DelayedAfun f
      -> LLVM arch (CompiledAfun arch f)
  compileAfun = compileOpenAfun
  
  
  {-# INLINEABLE compileOpenAfun #-}
  compileOpenAfun
      :: (HasCallStack, Compile arch)
      => DelayedOpenAfun aenv f
      -> LLVM arch (CompiledOpenAfun arch aenv f)
  compileOpenAfun (Alam lhs l) = Alam lhs <$> compileOpenAfun l
  compileOpenAfun (Abody b)    = Abody    <$> compileOpenAcc b
  
  
  {-# INLINEABLE compileOpenAcc #-}
  compileOpenAcc
      :: forall arch _aenv _a. (HasCallStack, Compile arch)
      => DelayedOpenAcc _aenv _a
      -> LLVM arch (CompiledOpenAcc arch _aenv _a)
  compileOpenAcc = traverseAcc
    where
      -- Traverse an open array expression in depth-first order. The top-level
      -- function 'traverseAcc' is intended for manifest arrays that we will
      -- generate LLVM code for. Array valued sub-terms, which might be manifest
      -- or delayed, are handled separately.
      --
      -- As the AST is traversed, we also collect a set of the indices of free
      -- array variables that were referred to within scalar sub-expressions.
      -- These will be required during code generation and execution.
      --
      traverseAcc
          :: forall aenv arrs. HasCallStack
          => DelayedOpenAcc aenv arrs
          -> LLVM arch (CompiledOpenAcc arch aenv arrs)
      traverseAcc Delayed{}       = internalError "unexpected delayed array"
      traverseAcc (Manifest pacc) =
        case pacc of
          -- Environment and control flow
          Avar ix                     -> plain $ pure (AST.Avar ix)
          Alet lhs a b                -> plain . pure =<< AST.Alet lhs  <$> traverseAcc a <*> traverseAcc b
          Apply r f a                 -> plain =<< liftA2 (AST.Apply r) <$> travAF f <*> travA a
          Awhile p f a                -> plain =<< liftA3 AST.Awhile    <$> travAF p <*> travAF f <*> travA a
          Acond p t e                 -> plain =<< liftA3 AST.Acond     <$> travE  p <*> travA  t <*> travA e
          Apair a1 a2                 -> plain =<< liftA2 AST.Apair     <$> travA a1 <*> travA a2
          Anil                        -> plain $ pure AST.Anil
          Atrace msg a1 a2            -> plain =<< liftA2 (AST.Atrace msg) <$> travA a1 <*> travA a2
  
          -- Foreign arrays operations
          Aforeign repr ff afun a     -> foreignA repr ff afun a
  
          -- Uninitialised array allocation
          Generate r sh f
            | alloc f                 -> plain =<< liftA (AST.Alloc r)  <$> travE sh
  
          -- Array injection & manipulation
          Reshape shr sh a            -> plain =<< liftA2 (AST.Reshape shr) <$> travE sh <*> travM a
          Unit tp e                   -> plain =<< liftA (AST.Unit tp)  <$> travE e
          Use repr arrs               -> plain $ pure (AST.Use repr arrs)
          Map _ f a
            | Just (t,x) <- unzip f a -> plain $ pure (AST.Unzip t x)
  
          -- Skeleton operations resulting in compiled code.
          -- Important: AST fragments that accelerate:DAA.Analysis.Hash ignores,
          -- must also be ignored here! Otherwise kernels with different aenvs
          -- get the same hash. These fragments are ignored using ignoreE.
  
          -- Producers
          Map tp f a                  -> build =<< liftA2 (map tp)      <$> travF f  <*> travA a
          Generate r sh f             -> build =<< liftA2 (generate r)  <$> ignoreE sh <*> travF f
          Transform r sh p f a        -> build =<< liftA4 (transform r) <$> ignoreE sh <*> travF p <*> travF f <*> travA a
          Backpermute shr sh f a      -> build =<< liftA3 (backpermute shr) <$> ignoreE sh <*> travF f <*> travA a
  
          -- Consumers
          Fold f z a                  -> build =<< liftA3 fold          <$> travF f <*> travME z <*> travD a
          FoldSeg i f z a s           -> build =<< liftA4 (foldSeg i)   <$> travF f <*> travME z <*> travD a <*> travD s
          Scan  d f z a               -> build =<< liftA3 (scan  d)     <$> travF f <*> travME z <*> travD a
          Scan' d f z a               -> build =<< liftA3 (scan' d)     <$> travF f <*> travE z <*> travD a
          Permute f d g a             -> build =<< liftA4 permute       <$> travF f <*> travA d <*> travF g <*> travD a
          Stencil s tp f x a          -> build =<< liftA3 (stencil1 s tp) <$> travF f <*> travB x <*> travD a
          Stencil2 s1 s2 tp f x a y b -> build =<< liftA5 (stencil2 s1 s2 tp) <$> travF f <*> travB x <*> travD a <*> travB y <*> travD b
  
          -- Removed by fusion
          Replicate{}                 -> fusionError
          Slice{}                     -> fusionError
          ZipWith{}                   -> fusionError
  
        where
          map tp _ a             = AST.Map tp a
          generate r sh _        = AST.Generate r sh
          transform r sh _ _ a   = AST.Transform r sh a
          backpermute shr sh _ a = AST.Backpermute shr sh a
          fold _ z a             = AST.Fold z a
          foldSeg i _ z a s      = AST.FoldSeg i z a s
          scan d _ z a           = AST.Scan d z a
          scan' d _ _ a          = AST.Scan' d a
          permute _ d _ a        = AST.Permute d a
  
          stencil1 :: StencilR sh a stencil
                   -> TypeR b
                   -> Fun      aenv (stencil -> b)
                   -> Boundary aenv (Array sh a)
                   -> AST.DelayedOpenAcc     CompiledOpenAcc arch aenv (Array sh a)
                   -> AST.PreOpenAccSkeleton CompiledOpenAcc arch aenv (Array sh b)
          stencil1 s tp _ _ a = AST.Stencil1 tp (snd $ stencilHalo s) a
  
          stencil2 :: StencilR sh a stencil1
                   -> StencilR sh b stencil2
                   -> TypeR c
                   -> Fun                                          aenv (stencil1 -> stencil2 -> c)
                   -> Boundary                                     aenv (Array sh a)
                   -> AST.DelayedOpenAcc     CompiledOpenAcc arch  aenv (Array sh a)
                   -> Boundary                                     aenv (Array sh b)
                   -> AST.DelayedOpenAcc     CompiledOpenAcc arch  aenv (Array sh b)
                   -> AST.PreOpenAccSkeleton CompiledOpenAcc arch  aenv (Array sh c)
          stencil2 s1 s2 tp _ _ a _ b = AST.Stencil2 tp (union shr h1 h2) a b
            where
              (shr, h1) = stencilHalo s1
              (_,   h2) = stencilHalo s2
  
          fusionError :: error
          fusionError = internalError ("unexpected fusible material: " % formatPreAccOp) pacc
  
          travA :: HasCallStack
                => DelayedOpenAcc aenv a
                -> LLVM arch (IntMap (Idx' aenv), CompiledOpenAcc arch aenv a)
          travA acc = pure <$> traverseAcc acc
  
          travD :: HasCallStack
                => DelayedOpenAcc aenv (Array sh e)
                -> LLVM arch ( IntMap (Idx' aenv)
                             , AST.DelayedOpenAcc CompiledOpenAcc arch aenv (Array sh e))
          travD acc =
            case acc of
              Delayed{..} -> liftA2 (const . (AST.Delayed reprD)) <$> travE extentD <*> travF indexD
              _           -> liftA  (AST.Manifest $ arraysR acc)  <$> travA acc
  
          travM :: HasCallStack
                => DelayedOpenAcc aenv (Array sh e)
                -> LLVM arch (IntMap (Idx' aenv), ArrayVar aenv (Array sh e))
          travM (Manifest (Avar v)) = return (freevar v, v)
          travM _                   = internalError "expected array variable"
  
          travME :: Maybe (OpenExp env aenv e) -> LLVM arch (IntMap (Idx' aenv), Bool)
          travME Nothing  = return (IntMap.empty, False)
          travME (Just e) = (True <$) <$> travE e
  
          travAF :: HasCallStack
                 => DelayedOpenAfun aenv f
                 -> LLVM arch (IntMap (Idx' aenv), CompiledOpenAfun arch aenv f)
          travAF afun = pure <$> compileOpenAfun afun
  
          travF :: HasCallStack
                => OpenFun env aenv t
                -> LLVM arch (IntMap (Idx' aenv), OpenFun env aenv t)
          travF (Body b)    = liftA Body <$> travE b
          travF (Lam lhs f) = liftA (Lam lhs) <$> travF f
  
          travB :: HasCallStack
                => Boundary aenv t
                -> LLVM arch (IntMap (Idx' aenv), Boundary aenv t)
          travB Clamp        = return $ pure Clamp
          travB Mirror       = return $ pure Mirror
          travB Wrap         = return $ pure Wrap
          travB (Constant c) = return $ pure (Constant c)
          travB (Function f) = liftA Function <$> travF f
  
          ignoreE :: HasCallStack
                  => OpenExp env aenv e
                  -> LLVM arch (IntMap (Idx' aenv), OpenExp env aenv e)
          ignoreE exp = return $ pure exp
  
          build :: (IntMap (Idx' aenv), AST.PreOpenAccSkeleton CompiledOpenAcc arch aenv arrs)
                -> LLVM arch (CompiledOpenAcc arch aenv arrs)
          build (aenv, eacc) = do
            let aval = makeGamma aenv
            kernel <- compileForTarget pacc aval
            return $! BuildAcc (arraysR eacc) aval kernel eacc
  
          plain :: (IntMap (Idx' aenv'), AST.PreOpenAccCommand CompiledOpenAcc arch aenv' arrs')
                -> LLVM arch (CompiledOpenAcc arch aenv' arrs')
          plain (_, eacc) = return (PlainAcc (arraysR eacc) eacc)
  
          -- Filling an array with undefined values is equivalent to allocating an
          -- uninitialised array. We look for this specific pattern because we
          -- expect it to appear only in certain places, e.g. as the default array
          -- in a 'permute' where the default values are never used. Note however
          -- that the simplifier does not take into account 'undef' values. For
          -- example, the following transformation is valid:
          --
          --   x + undef  ~~>  undef
          --
          -- so it is still possible to generate empty functions which we will
          -- execute, even though they do nothing (except incur scheduler
          -- overhead).
          --
          alloc :: Fun aenv (sh -> e)
                -> Bool
          alloc f
            | Lam _ (Body (Undef _)) <- f = True
            | otherwise                   = False
  
          -- Unzips of manifest array data can be done in constant time without
          -- executing any array programs. We split them out here into a separate
          -- case so that the execution phase does not have to continually perform
          -- the below check.
          --
          unzip :: forall sh a b.
                   Fun aenv (a -> b)
                -> DelayedOpenAcc aenv (Array sh a)
                -> Maybe (AST.UnzipIdx a b, ArrayVar aenv (Array sh a))
          unzip f a
            | Lam lhs (Body b) <- f
            , Just vars <- extractExpVars b
            , Delayed _ sh index _                <- a
            , Shape u                             <- sh
            , Just v                              <- isIdentityIndexing index
            , Just Refl                           <- matchVar u v
            = Just (unzipIdx lhs vars, u)
          unzip _ _
            = Nothing
  
          unzipIdx :: forall env a b. ELeftHandSide a () env -> Vars ScalarType env b -> AST.UnzipIdx a b
          unzipIdx lhs = go
            where
              go :: Vars ScalarType env y -> AST.UnzipIdx a y
              go TupRunit                = AST.UnzipUnit
              go (TupRpair v1 v2)        = AST.UnzipPair (go v1) (go v2)
              go (TupRsingle (Var _ ix)) = case lookupVar lhs ix of
                Right u -> u
                Left _  -> internalError "Left branch is unreachable as `Idx () y` is an empty type"
  
              lookupVar :: ELeftHandSide x env1 env2 -> Idx env2 y -> Either (Idx env1 y) (AST.UnzipIdx x y)
              lookupVar (LeftHandSideWildcard _) ix = Left ix
              lookupVar (LeftHandSideSingle _)   ix = case ix of
                ZeroIdx     -> Right AST.UnzipId
                SuccIdx ix' -> Left ix'
              lookupVar (LeftHandSidePair l1 l2) ix = case lookupVar l2 ix of
                Right u -> Right $ AST.UnzipPrj PairIdxRight u
                Left ix' -> case lookupVar l1 ix' of
                  Right u   -> Right $ AST.UnzipPrj PairIdxLeft u
                  Left ix'' -> Left ix''
  
          -- Is there a foreign version available for this backend? If so, take
          -- the foreign function and drop the remaining terms. Otherwise, drop
          -- this term and continue walking down the list of alternate
          -- implementations.
          --
          foreignA :: (HasCallStack, A.Foreign asm)
                   => ArraysR b
                   -> asm         (a -> b)
                   -> DelayedAfun (a -> b)
                   -> DelayedOpenAcc aenv a
                   -> LLVM arch (CompiledOpenAcc arch aenv b)
          foreignA repr ff f a =
            case foreignAcc ff of
              Just asm -> plain =<< liftA (AST.Aforeign repr (A.strForeign ff) asm) <$> travA a
              Nothing  -> traverseAcc $ Manifest (Apply repr (weaken weakenEmpty f) a)
  
      -- Traverse a scalar expression
      --
      travE :: HasCallStack
            => OpenExp env aenv e
            -> LLVM arch (IntMap (Idx' aenv), OpenExp env aenv e)
      travE exp =
        case exp of
          Evar v                  -> return $ pure $ Evar v
          Const tp c              -> return $ pure $ Const tp c
          PrimConst c             -> return $ pure $ PrimConst c
          Undef tp                -> return $ pure $ Undef tp
          Foreign tp ff f x       -> foreignE tp ff f x
          --
          Let lhs a b             -> liftA2 (Let lhs)         <$> travE a <*> travE b
          IndexSlice slix x s     -> liftA2 (IndexSlice slix) <$> travE x <*> travE s
          IndexFull slix x s      -> liftA2 (IndexFull slix)  <$> travE x <*> travE s
          ToIndex shr s i         -> liftA2 (ToIndex   shr)   <$> travE s <*> travE i
          FromIndex shr s i       -> liftA2 (FromIndex shr)   <$> travE s <*> travE i
          Nil                     -> return $ pure Nil
          Pair e1 e2              -> liftA2 Pair              <$> travE e1 <*> travE e2
          VecPack   vecr e        -> liftA  (VecPack   vecr)  <$> travE e
          VecUnpack vecr e        -> liftA  (VecUnpack vecr)  <$> travE e
          Case t xs x             -> liftA3 Case              <$> travE t <*> travLE xs <*> travME x
          Cond p t e              -> liftA3 Cond              <$> travE p <*> travE t <*> travE e
          While p f x             -> liftA3 While             <$> travF p <*> travF f <*> travE x
          PrimApp f e             -> liftA  (PrimApp f)       <$> travE e
          Index a e               -> liftA2 Index             <$> travA a <*> travE e
          LinearIndex a e         -> liftA2 LinearIndex       <$> travA a <*> travE e
          Shape a                 -> liftA  Shape             <$> travA a
          ShapeSize shr e         -> liftA  (ShapeSize shr)   <$> travE e
          Coerce t1 t2 x          -> liftA  (Coerce t1 t2)    <$> travE x
  
        where
          travA :: ArrayVar aenv (Array sh e)
                -> LLVM arch (IntMap (Idx' aenv), ArrayVar aenv (Array sh e))
          travA var = return (freevar var, var)
  
          travF :: HasCallStack
                => OpenFun env aenv t
                -> LLVM arch (IntMap (Idx' aenv), OpenFun env aenv t)
          travF (Body b)    = liftA Body      <$> travE b
          travF (Lam lhs f) = liftA (Lam lhs) <$> travF f
  
          travLE :: HasCallStack
                 => [(TAG, OpenExp env aenv t)]
                 -> LLVM arch (IntMap (Idx' aenv), [(TAG, OpenExp env aenv t)])
          travLE []     = return $ pure []
          travLE ((t,x):xs) = do
            (v,  y)  <- travE x
            (vs, ys) <- travLE xs
            return (IntMap.union v vs, (t,y):ys)
  
          travME :: HasCallStack
                 => Maybe (OpenExp env aenv t)
                 -> LLVM arch (IntMap (Idx' aenv), Maybe (OpenExp env aenv t))
          travME Nothing  = return $ pure Nothing
          travME (Just e) = fmap Just <$> travE e
  
          foreignE :: (HasCallStack, A.Foreign asm)
                   => TypeR b
                   -> asm           (a -> b)
                   -> Fun () (a -> b)
                   -> OpenExp env aenv a
                   -> LLVM arch (IntMap (Idx' aenv), OpenExp env aenv b)
          foreignE tp asm f x =
            case foreignExp @arch asm of
              Just{}                            -> liftA (Foreign tp asm err) <$> travE x
              Nothing
                | Lam lhs (Body b) <- f
                , Exists lhs' <- rebuildLHS lhs -> liftA2 (Let lhs')       <$> travE x <*> travE (weaken weakenEmpty $ weakenE (sinkWithLHS lhs lhs' weakenEmpty) b)
              _                                 -> error "the slow regard of silent things"
            where
              err :: Fun () (a -> b)
              err = internalError "attempt to use fallback in foreign expression"
  
  
  -- Applicative
  -- -----------
  --
  liftA4 :: Applicative f => (a -> b -> c -> d -> e) -> f a -> f b -> f c -> f d -> f e
  liftA4 f a b c d = f <$> a <*> b <*> c <*> d
  
  liftA5 :: Applicative f => (a -> b -> c -> d -> e -> g) -> f a -> f b -> f c -> f d -> f e -> f g
  liftA5 f a b c d g = f <$> a <*> b <*> c <*> d <*> g
  
  instance HasArraysR (CompiledOpenAcc arch) where
    {-# INLINEABLE arraysR #-}
    arraysR (BuildAcc r _ _ _) = r
    arraysR (PlainAcc r     _) = r
  
  