content: {-# LANGUAGE CPP                   #-}
  {-# LANGUAGE FlexibleInstances     #-}
  {-# LANGUAGE GADTs                 #-}
  {-# LANGUAGE ImpredicativeTypes    #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE PatternGuards         #-}
  {-# LANGUAGE QuasiQuotes           #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  {-# LANGUAGE TemplateHaskell       #-}
  #if __GLASGOW_HASKELL__ <= 708
  {-# LANGUAGE OverlappingInstances #-}
  {-# OPTIONS_GHC -fno-warn-unrecognised-pragmas #-}
  #endif
  
  -- |
  -- Module      : Data.Array.Accelerate.CUDA.CodeGen.Base
  -- Copyright   : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
  --               [2009..2014] Trevor L. McDonell
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  
  module Data.Array.Accelerate.CUDA.CodeGen.Base (
  
    -- Names and Types
    CUTranslSkel(..), CUDelayedAcc(..), CUExp(..), CUFun1(..), CUFun2(..),
    Eliminate, Instantiate1, Instantiate2,
    Name, namesOfArray, namesOfAvar, groupOfInt,
  
    -- Declaration generation
    cint, cvar, ccall, cchar, cintegral, cbool, cshape, cslice, csize, cindexHead, cindexTail, ctoIndex, cfromIndex,
    readArray, writeArray, shared,
    indexArray, environment, arrayAsTex, arrayAsArg,
    umul24, gridSize, threadIdx,
  
    -- Mutable operations
    (.=.), locals, Lvalue(..), Rvalue(..),
  
  ) where
  
  -- library
  import Prelude                                          hiding ( zipWith, zipWith3 )
  import Data.List                                        ( mapAccumR )
  import Text.PrettyPrint.Mainland
  import Language.C.Quote.CUDA
  import qualified Language.C.Syntax                      as C
  import qualified Data.HashMap.Strict                    as Map
  
  -- cuda
  import Foreign.CUDA.Analysis.Device
  
  -- friends
  import Data.Array.Accelerate.Type
  import Data.Array.Accelerate.Error
  import Data.Array.Accelerate.Array.Representation       ( SliceIndex(..) )
  import Data.Array.Accelerate.Array.Sugar                ( Array, Shape, Elt )
  import Data.Array.Accelerate.Analysis.Shape
  import Data.Array.Accelerate.CUDA.CodeGen.Type
  import Data.Array.Accelerate.CUDA.AST
  
  -- Names
  -- -----
  
  type Name = String
  
  namesOfArray
      :: forall e. Elt e
      => Name             -- name of group: typically "Out" or "InX" for some number 'X'
      -> e                -- dummy
      -> (Name, [Name])   -- shape and array field names
  namesOfArray grp _
    = let ty      = eltType (undefined :: e)
          arr x   = "arr" ++ grp ++ '_':show x
          n       = length ty
      in
      ( "sh" ++ grp, map arr [n-1, n-2 .. 0] )
  
  
  namesOfAvar :: forall aenv sh e. (Shape sh, Elt e) => Gamma aenv -> Idx aenv (Array sh e) -> (Name, [Name])
  namesOfAvar gamma ix = namesOfArray (groupOfAvar gamma ix) (undefined::e)
  
  groupOfAvar :: (Shape sh, Elt e) => Gamma aenv -> Idx aenv (Array sh e) -> Name
  groupOfAvar (Gamma gamma) = groupOfInt . (gamma Map.!) . Idx_
  
  groupOfInt :: Int -> Name
  groupOfInt n = "In" ++ show n
  
  
  -- Types of compilation units
  -- --------------------------
  
  -- A CUDA compilation unit, together with the name of the main __global__ entry
  -- function.
  --
  data CUTranslSkel aenv a = CUTranslSkel Name [C.Definition]
  
  instance Show (CUTranslSkel aenv a) where
    show (CUTranslSkel entry _) = entry
  
  instance Pretty (CUTranslSkel aenv a) where
    ppr  (CUTranslSkel _ code)  = ppr code
  
  -- Scalar expressions, including the environment of local let-bindings to bring
  -- into scope before evaluating the body.
  --
  data CUExp aenv a where
    CUExp  :: ([C.BlockItem], [C.Exp])
           -> CUExp aenv a
  
  -- Scalar functions of particular arity, with local bindings.
  --
  type Eliminate a        = forall x. [x] -> [(Bool,x)]
  type Instantiate1 a b   = forall x. Rvalue x => [x] -> ([C.BlockItem], [C.Exp])
  type Instantiate2 a b c = forall x y. (Rvalue x, Rvalue y) => [x] -> [y] -> ([C.BlockItem], [C.Exp])
  
  data CUFun1 aenv f where
    CUFun1 :: (Elt a, Elt b)
           => Eliminate a
           -> Instantiate1 a b
           -> CUFun1 aenv (a -> b)
  
  data CUFun2 aenv f where
    CUFun2 :: (Elt a, Elt b, Elt c)
           => Eliminate a
           -> Eliminate b
           -> Instantiate2 a b c
           -> CUFun2 aenv (a -> b -> c)
  
  -- Delayed arrays
  --
  data CUDelayedAcc aenv sh e where
    CUDelayed :: CUExp  aenv sh
              -> CUFun1 aenv (sh -> e)
              -> CUFun1 aenv (Int -> e)
              -> CUDelayedAcc aenv sh e
  
  
  -- Common expression forms
  -- -----------------------
  
  cint :: C.Type
  cint = typeOf (scalarType :: ScalarType Int)
  
  cvar :: Name -> C.Exp
  cvar x = [cexp|$id:x|]
  
  ccall :: Name -> [C.Exp] -> C.Exp
  ccall fn args = [cexp|$id:fn ($args:args)|]
  
  cchar :: Char -> C.Exp
  cchar c = [cexp|$char:c|]
  
  cintegral :: (Integral a, Show a) => a -> C.Exp
  cintegral n = [cexp|$int:n|]
  
  cbool :: Bool -> C.Exp
  cbool = cintegral . fromEnum
  
  cslice :: SliceIndex slix sl co dim -> Name -> ([C.Param], [C.Exp], [(C.Type, Name)])
  cslice slix sl =
    let xs = cshape' (ncodims slix) sl
        args = [ [cparam| const $ty:cint $id:x |] | x <- xs ]
    in (args, map cvar xs, zip (repeat cint) xs)
    where
      ncodims :: SliceIndex slix sl co dim -> Int
      ncodims SliceNil = 0
      ncodims (SliceAll   s) = ncodims s
      ncodims (SliceFixed s) = ncodims s + 1
  
  -- Generate all the names of a shape given a base name and dimensionality
  cshape :: Int -> Name -> [C.Exp]
  cshape dim sh = [ cvar x | x <- cshape' dim sh ]
  
  cshape' :: Int -> Name -> [Name]
  cshape' dim sh = [ (sh ++ '_':show i) | i <- [dim-1, dim-2 .. 0] ]
  
  -- Get the innermost index of a shape/index
  cindexHead :: Rvalue r => [r] -> C.Exp
  cindexHead = rvalue . last
  
  -- Get the tail of a shape/index
  cindexTail :: Rvalue r => [r] -> [C.Exp]
  cindexTail = map rvalue . init
  
  -- generate code that calculates the product of the list of expressions
  csize :: Rvalue r => [r] -> C.Exp
  csize [] = [cexp| 1 |]
  csize ss = foldr1 (\a b -> [cexp| $exp:a * $exp:b |]) (map rvalue ss)
  
  -- Generate code to calculate a linear from a multi-dimensional index (given an array shape).
  --
  ctoIndex :: (Rvalue sh, Rvalue ix) => [sh] -> [ix] -> C.Exp
  ctoIndex extent index
    = toIndex (reverse $ map rvalue extent) (reverse $ map rvalue index)  -- we use a row-major representation
    where
      toIndex []      []     = [cexp| $int:(0::Int) |]
      toIndex [_]     [i]    = i
      toIndex (sz:sh) (i:ix) = [cexp| $exp:(toIndex sh ix) * $exp:sz + $exp:i |]
      toIndex _       _      = $internalError "toIndex" "argument mismatch"
  
  -- Generate code to calculate a multi-dimensional index from a linear index and a given array shape.
  -- This version creates temporary values that are reused in the computation.
  --
  cfromIndex :: (Rvalue sh, Rvalue ix) => [sh] -> ix -> Name -> ([C.BlockItem], [C.Exp])
  cfromIndex shName ixName tmpName = fromIndex (map rvalue shName) (rvalue ixName)
    where
      fromIndex [sh]   ix = ([], [[cexp| ({ assert( $exp:ix >= 0 && $exp:ix < $exp:sh ); $exp:ix; }) |]])
      fromIndex extent ix = let ((env, _, _), sh) = mapAccumR go ([], ix, 0) extent
                            in  (reverse env, sh)
  
      go (tmps,ix,n) d
        = let tmp         = tmpName ++ '_':show (n::Int)
              ix'         = [citem| const $ty:cint $id:tmp = $exp:ix ; |]
          in
          ((ix':tmps, [cexp| $id:tmp / $exp:d |], n+1), [cexp| $id:tmp % $exp:d |])
  
  
  -- Thread blocks and indices
  -- -------------------------
  
  umul24 :: DeviceProperties -> C.Exp -> C.Exp -> C.Exp
  umul24 dev x y
    | computeCapability dev < Compute 2 0 = [cexp| __umul24($exp:x, $exp:y) |]
    | otherwise                           = [cexp| $exp:x * $exp:y |]
  
  gridSize :: DeviceProperties -> C.Exp
  gridSize dev = umul24 dev [cexp|blockDim.x|] [cexp|gridDim.x|]
  
  threadIdx :: DeviceProperties -> C.Exp
  threadIdx dev =
    let block = umul24 dev [cexp|blockDim.x|] [cexp|blockIdx.x|]
    in  [cexp| $exp:block + threadIdx.x |]
  
  
  -- Generate an array indexing expression. Depending on the hardware class, this
  -- will be via direct array indexing or texture references.
  --
  indexArray
      :: DeviceProperties
      -> C.Type                   -- array element type (Float, Double...)
      -> C.Exp                    -- array variable name (arrInX_Y)
      -> C.Exp                    -- linear index
      -> C.Exp
  indexArray dev elt arr ix
    -- use the L2 cache of newer devices
    | computeCapability dev >= Compute 2 0 = [cexp| $exp:arr [ $exp:ix ] |]
  
    -- use the texture cache of compute 1.x devices
    | [cty|double|] <- elt                 = ccall "indexDArray" [arr, ix]
    | otherwise                            = ccall "indexArray"  [arr, ix]
  
  
  -- Kernel function parameters
  -- --------------------------
  
  -- Generate kernel parameters for an array valued argument, and a function to
  -- linearly index this array. Note that dimensional indexing results in error.
  --
  readArray
      :: forall aenv sh e. (Shape sh, Elt e)
      => Name                             -- group names
      -> Array sh e                       -- dummy to fix types
      -> ( [C.Param]
         , [C.Exp]
         , CUDelayedAcc aenv sh e )
  readArray grp dummy
    = let (sh, arrs)      = namesOfArray grp (undefined :: e)
          args            = arrayAsArg dummy grp
  
          dim             = expDim (undefined :: Exp aenv sh)
          sh'             = cshape dim sh
          get ix          = ([], map (\a -> [cexp| $id:a [ $exp:ix ] |]) arrs)
          manifest        = CUDelayed (CUExp ([], sh'))
                                      ($internalError "readArray" "linear indexing only")
                                      (CUFun1 (zip (repeat True)) (\[i] -> get (rvalue i)))
      in ( args, sh', manifest )
  
  
  -- Generate function parameters and corresponding variable names for the
  -- components of the given output array. The parameter list generated is
  -- suitable for marshalling an instance of "Array sh e", consisting of a group
  -- name (say "Out") to be welded with a shape name "shOut" followed by the
  -- non-parametric array data "arrOut_aX".
  --
  writeArray
      :: forall sh e. (Shape sh, Elt e)
      => Name                                     -- group names
      -> Array sh e                               -- dummy to fix types
      -> ( [C.Param]                              -- function parameters to marshal the output array
         , [C.Exp]                                -- the shape of the output array
         , forall x. Rvalue x => x -> [C.Exp] )   -- write an element at a given index
  writeArray grp _ =
    let (sh, arrs)        = namesOfArray grp (undefined :: e)
        dim               = expDim (undefined :: Exp aenv sh)
        sh'               = cshape' dim sh
        extent            = [ [cparam| const $ty:cint $id:i |] | i <- sh' ]
        adata             = zipWith (\t n -> [cparam| $ty:t * __restrict__ $id:n |]) (eltType (undefined :: e)) arrs
    in
    ( extent ++ adata
    , map cvar sh'
    , \ix -> map (\a -> [cexp| $id:a [ $exp:(rvalue ix) ] |]) arrs
    )
  
  
  -- All dynamically allocated __shared__ memory will begin at the same base
  -- address. If we call this more than once, or the kernel itself declares some
  -- shared memory, the first parameter is a pointer to where the new declarations
  -- should take as the base address.
  --
  shared
      :: forall e. Elt e
      => e                                        -- dummy type
      -> Name                                     -- group name
      -> C.Exp                                    -- how much shared memory per type
      -> Maybe C.Exp                              -- (optional) initialise from this base address
      -> ( [C.InitGroup]                          -- shared memory declaration and...
         , forall x. Rvalue x => x -> [C.Exp])    -- ...indexing function
  shared _ grp size mprev
    = let e:es                    = eltType (undefined :: e)
          x:xs                    = let k = length es in map (\n -> grp ++ show n) [k, k-1 .. 0]
  
          sdata t v p             = [cdecl| volatile $ty:t * $id:v = ($ty:t *) & $id:p [ $exp:size ]; |]
          sbase t v
            | Just p <- mprev     = [cdecl| volatile $ty:t * $id:v = ($ty:t *) $exp:p; |]
            | otherwise           = [cdecl| extern volatile __shared__ $ty:t $id:v [] ; |]
      in
      ( sbase e x : zipWith3 sdata es xs (init (x:xs))
      , \ix -> map (\v -> [cexp| $id:v [ $exp:(rvalue ix) ] |]) (x:xs)
      )
  
  -- Array environment references. The method in which arrays are accessed depends
  -- on the device architecture (see below). We always include the array shape
  -- before the array data terms.
  --
  --   compute 1.x:
  --      texture references of type [Definition]
  --
  --   compute 2.x and 3.x:
  --      function arguments of type [Param]
  --
  -- NOTE: The environment variables must always be the first argument to the
  --       kernel function, as this is where they will be marshaled during the
  --       execution phase.
  --
  environment
      :: forall aenv. DeviceProperties
      -> Gamma aenv
      -> ([C.Definition], [C.Param])
  environment dev gamma@(Gamma aenv)
    | computeCapability dev < Compute 2 0
    = Map.foldrWithKey (\(Idx_ v) _ (ds,ps) -> let (d,p) = asTex v in (d++ds, p++ps)) ([],[]) aenv
  
    | otherwise
    = ([], Map.foldrWithKey (\(Idx_ v) _ vs -> asArg v ++ vs) [] aenv)
  
    where
      asTex :: forall sh e. (Shape sh, Elt e) => Idx aenv (Array sh e) -> ([C.Definition], [C.Param])
      asTex ix = arrayAsTex (undefined :: Array sh e) (groupOfAvar gamma ix)
  
      asArg :: forall sh e. (Shape sh, Elt e) => Idx aenv (Array sh e) -> [C.Param]
      asArg ix = arrayAsArg (undefined :: Array sh e) (groupOfAvar gamma ix)
  
  
  arrayAsTex :: forall sh e. (Shape sh, Elt e) => Array sh e -> Name -> ([C.Definition], [C.Param])
  arrayAsTex _ grp =
    let (sh, arrs)        = namesOfArray grp (undefined :: e)
        dim               = expDim (undefined :: Exp aenv sh)
        extent            = [ [cparam| const $ty:cint $id:i |] | i <- cshape' dim sh ]
        adata             = zipWith (\t a -> [cedecl| static $ty:t $id:a; |]) (eltTypeTex (undefined :: e)) arrs
    in
    (adata, extent)
  
  arrayAsArg :: forall sh e. (Shape sh, Elt e) => Array sh e -> Name -> [C.Param]
  arrayAsArg _ grp =
    let (sh, arrs)        = namesOfArray grp (undefined :: e)
        dim               = expDim (undefined :: Exp aenv sh)
        extent            = [ [cparam| const $ty:cint $id:i |] | i <- cshape' dim sh ]
        adata             = zipWith (\t n -> [cparam| const $ty:t * __restrict__ $id:n |]) (eltType (undefined :: e)) arrs
    in
    extent ++ adata
  
  
  -- Mutable operations
  -- ------------------
  
  -- Declare some local variables. These can be either const or mutable
  -- declarations.
  --
  locals :: forall e. Elt e
         => Name
         -> e
         -> ( [(C.Type, Name)]            -- const declarations
            , [C.Exp], [C.InitGroup])     -- mutable declaration and names
  locals base _
    = let elt             = eltType (undefined :: e)
          n               = length elt
          local t v       = let name = base ++ show v
                            in ( (t, name), cvar name, [cdecl| $ty:t $id:name; |] )
      in
      unzip3 $ zipWith local elt [n-1, n-2 .. 0]
  
  class Lvalue a where
    lvalue :: a -> C.Exp -> C.BlockItem
  
  -- Note: [Mutable l-values]
  --
  -- Be careful when using mutable l-values that the same variable does not appear
  -- on both the left and right hand side. For example, the following will lead to
  -- problems (#114, #168):
  --
  --   $items:(x .=. f x)
  --
  --   $items:(y .=. combine x y)
  --
  -- If 'x' and 'y' represent values with tuple types, they will have multiple
  -- components. Since the LHS is updated as the new values are calculated, it is
  -- possible to get into a situation where computing the new value for some of
  -- the components will be using the updated values, not the original values.
  --
  -- Instead, store the result to some (temporary) variable that does not appear
  -- on the RHS, and then update old value using the fully computed result, e.g.:
  --
  --   $items:(x' .=. f x)
  --   $items:(x  .=. x')
  --
  instance Lvalue C.Exp where
    lvalue x y = [citem| $exp:x = $exp:y; |]
  
  instance Lvalue (C.Type, Name) where
    lvalue (t,x) y = [citem| const $ty:t $id:x = $exp:y; |]
  
  
  class Rvalue a where
    rvalue :: a -> C.Exp
  
  instance Rvalue C.Exp where
    rvalue = id
  
  instance Rvalue Name where
    rvalue = cvar
  
  instance Rvalue (C.Type, Name) where
    rvalue (_,x) = rvalue x
  
  
  infixr 0 .=.
  (.=.) :: Assign l r => l -> r -> [C.BlockItem]
  (.=.) =  assign
  
  class Assign l r where
    assign :: l -> r -> [C.BlockItem]
  
  instance (Lvalue l, Rvalue r) => Assign l r where
    assign lhs rhs = [ lvalue lhs (rvalue rhs) ]
  
  instance {-# OVERLAPS #-} Assign l r => Assign (Bool,l) r where
    assign (used,lhs) rhs
      | used      = assign lhs rhs
      | otherwise = []
  
  instance {-# OVERLAPS #-} Assign l r => Assign [l] [r] where
    assign []     []     = []
    assign (x:xs) (y:ys) = assign x y ++ assign xs ys
    assign _      _      = $internalError ".=." "argument mismatch"
  
  instance {-# OVERLAPS #-} Assign l r => Assign l ([C.BlockItem], r) where
    assign lhs (env, rhs) = env ++ assign lhs rhs
  
  
  -- Prelude'
  -- --------
  
  -- A version of zipWith that requires the lists to be equal length
  --
  zipWith :: (a -> b -> c) -> [a] -> [b] -> [c]
  zipWith f (x:xs) (y:ys) = f x y : zipWith f xs ys
  zipWith _ []     []     = []
  zipWith _ _      _      = $internalError "zipWith" "argument mismatch"
  
  zipWith3 :: (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
  zipWith3 f (x:xs) (y:ys) (z:zs) = f x y z : zipWith3 f xs ys zs
  zipWith3 _ []     []     []     = []
  zipWith3 _ _      _      _      = $internalError "zipWith3" "argument mismatch"
  
  