{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE TypeSynonymInstances #-} -- | -- Module : Data.Array.Accelerate.CUDA.AST -- Copyright : [2008..2014] Manuel M T Chakravarty, Gabriele Keller -- [2009..2014] Trevor L. McDonell -- License : BSD3 -- -- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- module Data.Array.Accelerate.CUDA.AST ( module Data.Array.Accelerate.AST, AccKernel(..), Free, Gamma(..), Idx_(..), ExecAcc, ExecAfun, ExecOpenAfun, ExecOpenAcc(..), ExecExp, ExecFun, ExecOpenExp, ExecOpenFun, -- ExecSeq(..), ExecOpenSeq(..), ExecP(..), ExecC(..), freevar, makeEnvMap, ) where -- friends import Data.Array.Accelerate.AST import Data.Array.Accelerate.Lifetime import Data.Array.Accelerate.Pretty as PP import Data.Array.Accelerate.Array.Sugar ( Array, Shape, Elt ) import qualified Data.Array.Accelerate.FullList as FL import qualified Foreign.CUDA.Driver as CUDA import qualified Foreign.CUDA.Analysis as CUDA -- system import Text.PrettyPrint import Data.Hashable import Data.Monoid hiding ( (<>) ) import qualified Data.HashSet as Set import qualified Data.HashMap.Strict as Map import Prelude -- A non-empty list of binary objects will be used to execute a kernel. We keep -- auxiliary information together with the compiled module, such as entry point -- and execution information. -- data AccKernel a where AccKernel :: !String -- __global__ entry function name -> {-# UNPACK #-} !CUDA.Fun -- __global__ function object -> {-# UNPACK #-} !(Lifetime CUDA.Module) -- binary module -> {-# UNPACK #-} !CUDA.Occupancy -- occupancy analysis -> {-# UNPACK #-} !Int -- thread block size -> {-# UNPACK #-} !Int -- shared memory per block (bytes) -> !(Int -> Int) -- number of blocks for input problem size -> AccKernel a -- Kernel execution is asynchronous, barriers allow (cross-stream) -- synchronisation to determine when the operation has completed -- -- data AccBarrier = AB !Stream !Event -- The set of free array variables for array computations that were embedded -- within scalar expressions. These arrays are are required to execute the -- kernel, by binding to texture references to similar. -- type Free aenv = Set.HashSet (Idx_ aenv) freevar :: (Shape sh, Elt e) => Idx aenv (Array sh e) -> Free aenv freevar = Set.singleton . Idx_ -- A mapping between environment indexes and some token identifying that array -- in the generated code. This simply compresses the sequence of array indices -- into a continuous range, rather than directly using the integer equivalent of -- the de Bruijn index. -- -- This results in generated code that is (slightly) less sensitive to the -- placement of let bindings, ultimately leading to a higher hit rate in the -- compilation cache. -- newtype Gamma aenv = Gamma ( Map.HashMap (Idx_ aenv) Int ) deriving ( Monoid ) makeEnvMap :: Free aenv -> Gamma aenv makeEnvMap indices = Gamma . Map.fromList . flip zip [0..] -- . sortBy (compare `on` idxType) $ Set.toList indices -- where -- idxType :: Idx_ aenv -> TypeRep -- idxType (Idx_ (_ :: Idx aenv (Array sh e))) = typeOf (undefined :: e) -- Opaque array environment indices -- data Idx_ aenv where Idx_ :: (Shape sh, Elt e) => Idx aenv (Array sh e) -> Idx_ aenv instance Eq (Idx_ aenv) where Idx_ ix1 == Idx_ ix2 = idxToInt ix1 == idxToInt ix2 instance Hashable (Idx_ aenv) where hashWithSalt salt (Idx_ ix) = salt `hashWithSalt` idxToInt ix -- Interleave compilation & execution state annotations into an open array -- computation AST -- data ExecOpenAcc aenv a where ExecAcc :: {-# UNPACK #-} !(FL.FullList () (AccKernel a)) -- executable binary objects -> !(Gamma aenv) -- free array variables the kernel needs access to -> !(PreOpenAcc ExecOpenAcc aenv a) -- the actual computation -> ExecOpenAcc aenv a -- the recursive knot EmbedAcc :: (Shape sh, Elt e) => !(PreExp ExecOpenAcc aenv sh) -- shape of the result array, used by execution -> ExecOpenAcc aenv (Array sh e) -- ExecSeq :: Arrays arrs -- => ExecOpenSeq aenv () arrs -- -> ExecOpenAcc aenv arrs -- An annotated AST suitable for execution in the CUDA environment -- type ExecAcc a = ExecOpenAcc () a type ExecAfun a = PreAfun ExecOpenAcc a type ExecOpenAfun aenv a = PreOpenAfun ExecOpenAcc aenv a type ExecOpenExp = PreOpenExp ExecOpenAcc type ExecOpenFun = PreOpenFun ExecOpenAcc type ExecExp = ExecOpenExp () type ExecFun = ExecOpenFun () -- Display the annotated AST -- ------------------------- instance Show (ExecAcc a) where show = render . prettyExecAcc noParens PP.Empty instance Show (ExecAfun a) where show = render . prettyExecAfun prettyExecAfun :: ExecAfun a -> Doc prettyExecAfun pfun = prettyPreOpenAfun prettyExecAcc PP.Empty pfun prettyExecAcc :: PrettyAcc ExecOpenAcc prettyExecAcc wrap aenv exec = case exec of EmbedAcc sh -> wrap $ hang (text "Embedded") 2 $ sep [ prettyPreExp prettyExecAcc parens aenv sh ] ExecAcc _ (Gamma fv) pacc -> let base = prettyPreOpenAcc prettyExecAcc wrap aenv pacc ann = braces (freevars (Map.keys fv)) freevars = (text "fv=" <>) . brackets . hcat . punctuate comma . map (\(Idx_ ix) -> char 'a' <> int (idxToInt ix)) in case pacc of Avar{} -> base Alet{} -> base Apply{} -> base Acond{} -> base Atuple{} -> base Aprj{} -> base _ -> ann <+> base -- ExecSeq _ -> text "" {-- data ExecSeq a where ExecS :: Extend ExecOpenAcc () aenv -> ExecOpenSeq aenv () a -> ExecSeq a data ExecOpenSeq aenv lenv arrs where ExecP :: Arrays a => ExecP aenv lenv a -> ExecOpenSeq aenv (lenv, a) arrs -> ExecOpenSeq aenv lenv arrs ExecC :: (Arrays a) => ExecC aenv lenv a -> ExecOpenSeq aenv lenv a ExecR :: Idx lenv a -> Maybe a -> ExecOpenSeq aenv lenv [a] data ExecP aenv lenv a where ExecToSeq :: (Elt slix, Shape sl, Shape sh, Elt e) => SliceIndex (EltRepr slix) (EltRepr sl) co (EltRepr sh) -> ExecOpenAcc aenv (Array sh e) -> AccKernel (Array sl e) -> !(Gamma aenv) -> [slix] -> ExecP aenv lenv (Array sl e) ExecUseLazy :: (Elt slix, Shape sl, Shape sh, Elt e) => SliceIndex (EltRepr slix) (EltRepr sl) co (EltRepr sh) -> Array sh e -> [slix] -> ExecP aenv lenv (Array sl e) ExecStreamIn :: Arrays a => [a] -> ExecP aenv lenv a ExecMap :: (Arrays a, Arrays b) => ExecOpenAfun aenv (a -> b) -> Idx lenv a -> ExecP aenv lenv b ExecZipWith :: (Arrays a, Arrays b, Arrays c) => ExecOpenAfun aenv (a -> b -> c) -> Idx lenv a -> Idx lenv b -> ExecP aenv lenv c ExecScanSeq :: Elt a => ExecFun aenv (a -> a -> a) -> ExecExp aenv a -> Idx lenv (Scalar a) -> Maybe a -> ExecP aenv lenv (Scalar a) data ExecC aenv lenv a where ExecFoldSeq :: Elt a => ExecFun aenv (a -> a -> a) -> ExecExp aenv a -> Idx lenv (Scalar a) -> Maybe a -> ExecC aenv lenv (Scalar a) ExecFoldSeqFlatten :: (Arrays a, Shape sh, Elt e) => ExecOpenAfun aenv (a -> Vector sh -> Vector e -> a) -> ExecOpenAcc aenv a -> Idx lenv (Array sh e) -> Maybe a -> ExecC aenv lenv a ExecStuple :: (Arrays a, IsAtuple a) => Atuple (ExecC aenv senv) (TupleRepr a) -> ExecC aenv senv a --}