accelerate-llvm-1.3.0.0: src/Data/Array/Accelerate/LLVM/AST.hs
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.AST
-- 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.AST (
DelayedOpenAcc(..),
PreOpenAccCommand(..),
PreOpenAccSkeleton(..),
UnzipIdx(..),
HasInitialValue,
) where
import Data.Array.Accelerate.LLVM.Execute.Async
import Data.Array.Accelerate.AST ( PreOpenAfun(..), HasArraysR(..), ArrayVar, ALeftHandSide, Exp, Direction, PrimBool, arrayR )
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Type
-- | Non-computational array program operations, parameterised over array
-- variables represented as de Bruijn indices.
--
data PreOpenAccCommand acc arch aenv a where
Avar :: ArrayVar aenv arrs
-> PreOpenAccCommand acc arch aenv arrs
Alet :: ALeftHandSide bnd aenv aenv'
-> acc arch aenv bnd
-> acc arch aenv' body
-> PreOpenAccCommand acc arch aenv body
Alloc :: ArrayR (Array sh e)
-> Exp aenv sh
-> PreOpenAccCommand acc arch aenv (Array sh e)
Use :: ArrayR (Array sh e)
-> Array sh e
-> PreOpenAccCommand acc arch aenv (Array sh e)
Unit :: TypeR e
-> Exp aenv e
-> PreOpenAccCommand acc arch aenv (Scalar e)
Apair :: acc arch aenv arrs1
-> acc arch aenv arrs2
-> PreOpenAccCommand acc arch aenv (arrs1, arrs2)
Anil :: PreOpenAccCommand acc arch aenv ()
Apply :: ArraysR bs
-> PreOpenAfun (acc arch) aenv (as -> bs)
-> acc arch aenv as
-> PreOpenAccCommand acc arch aenv bs
Aforeign :: ArraysR bs
-> String
-> (as -> Par arch (FutureR arch bs))
-> acc arch aenv as
-> PreOpenAccCommand acc arch aenv bs
Acond :: Exp aenv PrimBool
-> acc arch aenv arrs
-> acc arch aenv arrs
-> PreOpenAccCommand acc arch aenv arrs
Awhile :: PreOpenAfun (acc arch) aenv (arrs -> Scalar PrimBool)
-> PreOpenAfun (acc arch) aenv (arrs -> arrs)
-> acc arch aenv arrs
-> PreOpenAccCommand acc arch aenv arrs
Reshape :: ShapeR sh
-> Exp aenv sh
-> ArrayVar aenv (Array sh' e)
-> PreOpenAccCommand acc arch aenv (Array sh e)
Unzip :: UnzipIdx tup e
-> ArrayVar aenv (Array sh tup)
-> PreOpenAccCommand acc arch aenv (Array sh e)
-- | Collective array computations parameterised over array variables
-- represented as de Bruijn indices.
--
data PreOpenAccSkeleton acc arch aenv a where
-- Producers. The only way these terms can appear in the AST is if they
-- are applied to a manifest array.
--
Map :: TypeR b
-> acc arch aenv (Array sh a)
-> PreOpenAccSkeleton acc arch aenv (Array sh b)
Generate :: ArrayR (Array sh e)
-> Exp aenv sh
-> PreOpenAccSkeleton acc arch aenv (Array sh e)
Transform :: ArrayR (Array sh' b)
-> Exp aenv sh'
-> acc arch aenv (Array sh a)
-> PreOpenAccSkeleton acc arch aenv (Array sh' b)
Backpermute :: ShapeR sh'
-> Exp aenv sh'
-> acc arch aenv (Array sh e)
-> PreOpenAccSkeleton acc arch aenv (Array sh' e)
-- Consumers. These may have been applied to either manifest or delayed
-- array data.
--
Fold :: HasInitialValue
-> DelayedOpenAcc acc arch aenv (Array (sh, Int) e)
-> PreOpenAccSkeleton acc arch aenv (Array sh e)
FoldSeg :: IntegralType i
-> HasInitialValue
-> DelayedOpenAcc acc arch aenv (Array (sh, Int) e)
-> DelayedOpenAcc acc arch aenv (Segments i)
-> PreOpenAccSkeleton acc arch aenv (Array (sh, Int) e)
Scan :: Direction
-> HasInitialValue
-> DelayedOpenAcc acc arch aenv (Array (sh, Int) e)
-> PreOpenAccSkeleton acc arch aenv (Array (sh, Int) e)
Scan' :: Direction
-> DelayedOpenAcc acc arch aenv (Array (sh, Int) e)
-> PreOpenAccSkeleton acc arch aenv (Array (sh, Int) e, Array sh e)
Permute :: acc arch aenv (Array sh' e) -- target array (default values)
-> DelayedOpenAcc acc arch aenv (Array sh e) -- source values
-> PreOpenAccSkeleton acc arch aenv (Array sh' e)
Stencil1 :: TypeR b
-> sh -- stencil offset/halo size
-> DelayedOpenAcc acc arch aenv (Array sh a)
-> PreOpenAccSkeleton acc arch aenv (Array sh b)
Stencil2 :: TypeR c
-> sh -- stencil offset/halo size
-> DelayedOpenAcc acc arch aenv (Array sh a)
-> DelayedOpenAcc acc arch aenv (Array sh b)
-> PreOpenAccSkeleton acc arch aenv (Array sh c)
data UnzipIdx a b where
UnzipId :: UnzipIdx a a
UnzipPrj :: PairIdx a b -> UnzipIdx b c -> UnzipIdx a c
UnzipUnit :: UnzipIdx a ()
UnzipPair :: UnzipIdx a b1 -> UnzipIdx a b2 -> UnzipIdx a (b1, b2)
-- Denotes whether the fold or scan has an initial value.
-- When False, this is a fold1 or scan1.
--
type HasInitialValue = Bool
-- | Representation for array arguments.
--
-- If the argument is a delayed array (that is, it was fused into its
-- consumer) we only need to keep track of the extent of the argument. If
-- the argument is a manifest array, we recurse into the subterm.
--
data DelayedOpenAcc acc arch aenv a where
Delayed :: ArrayR (Array sh e)
-> Exp aenv sh
-> DelayedOpenAcc acc arch aenv (Array sh e)
Manifest :: ArraysR (Array sh e)
-> acc arch aenv (Array sh e)
-> DelayedOpenAcc acc arch aenv (Array sh e)
instance HasArraysR (acc arch) => HasArraysR (PreOpenAccCommand acc arch) where
{-# INLINEABLE arraysR #-}
arraysR (Avar (Var repr _)) = TupRsingle repr
arraysR (Alet _ _ a) = arraysR a
arraysR (Alloc repr _) = TupRsingle repr
arraysR (Use repr _) = TupRsingle repr
arraysR (Unit tp _) = TupRsingle $ ArrayR ShapeRz tp
arraysR (Apair a1 a2) = arraysR a1 `TupRpair` arraysR a2
arraysR Anil = TupRunit
arraysR (Apply repr _ _) = repr
arraysR (Aforeign repr _ _ _) = repr
arraysR (Acond _ a1 _) = arraysR a1
arraysR (Awhile _ _ a) = arraysR a
arraysR (Reshape shr _ (Var (ArrayR _ tp) _)) = TupRsingle $ ArrayR shr tp
arraysR (Unzip idx (Var (ArrayR shr tp) _)) = TupRsingle $ ArrayR shr $ go idx tp
where
go :: UnzipIdx a b -> TypeR a -> TypeR b
go UnzipId t = t
go (UnzipPrj PairIdxLeft ix) (TupRpair t _) = go ix t
go (UnzipPrj PairIdxRight ix) (TupRpair _ t) = go ix t
go UnzipUnit _ = TupRunit
go (UnzipPair ix1 ix2) t = go ix1 t `TupRpair` go ix2 t
go _ _ = error "Time enough for life to unfold all the precious things life has in store."
instance HasArraysR (acc arch) => HasArraysR (PreOpenAccSkeleton acc arch) where
{-# INLINEABLE arraysR #-}
arraysR (Map tp a) = let ArrayR shr _ = arrayR a
in TupRsingle $ ArrayR shr tp
arraysR (Generate repr _) = TupRsingle repr
arraysR (Transform repr _ _) = TupRsingle repr
arraysR (Backpermute shr _ a) = TupRsingle $ ArrayR shr $ arrayRtype $ arrayR a
arraysR (Fold _ a) = let ArrayR (ShapeRsnoc shr) tp = arrayR a
in TupRsingle $ ArrayR shr tp
arraysR (FoldSeg _ _ a _) = arraysR a
arraysR (Scan _ _ a) = arraysR a
arraysR (Scan' _ a) = let ArrayR (ShapeRsnoc shr) tp = arrayR a
in TupRsingle (ArrayR (ShapeRsnoc shr) tp) `TupRpair` TupRsingle (ArrayR shr tp)
arraysR (Permute a _) = arraysR a
arraysR (Stencil1 tp _ a) = let ArrayR shr _ = arrayR a
in TupRsingle $ ArrayR shr tp
arraysR (Stencil2 tp _ a _) = let ArrayR shr _ = arrayR a
in TupRsingle $ ArrayR shr tp
instance HasArraysR (acc arch) => HasArraysR (DelayedOpenAcc acc arch) where
{-# INLINEABLE arraysR #-}
arraysR (Delayed repr _) = TupRsingle repr
arraysR (Manifest repr _) = repr