{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}

-- | A representation with flat parallelism via GPU-oriented kernels.
module Futhark.IR.GPU
  ( GPU,

    -- * Module re-exports
    module Futhark.IR.Prop,
    module Futhark.IR.Traversals,
    module Futhark.IR.Pretty,
    module Futhark.IR.Syntax,
    module Futhark.IR.GPU.Op,
    module Futhark.IR.GPU.Sizes,
    module Futhark.IR.SOACS.SOAC,
  )
where

import Futhark.Builder
import Futhark.Construct
import Futhark.IR.GPU.Op
import Futhark.IR.GPU.Sizes
import Futhark.IR.Pretty
import Futhark.IR.Prop
import Futhark.IR.SOACS.SOAC hiding (HistOp (..))
import Futhark.IR.Syntax
import Futhark.IR.Traversals
import qualified Futhark.IR.TypeCheck as TC

-- | The phantom data type for the kernels representation.
data GPU

instance RepTypes GPU where
  type Op GPU = HostOp GPU (SOAC GPU)

instance ASTRep GPU where
  expTypesFromPat :: Pat (LetDec GPU) -> m [BranchType GPU]
expTypesFromPat = [ExtType] -> m [ExtType]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExtType] -> m [ExtType])
-> (Pat Type -> [ExtType]) -> Pat Type -> m [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat Type -> [ExtType]
forall dec. Typed dec => Pat dec -> [ExtType]
expExtTypesFromPat

instance TC.CheckableOp GPU where
  checkOp :: OpWithAliases (Op GPU) -> TypeM GPU ()
checkOp = Maybe SegLevel
-> HostOp (Aliases GPU) (SOAC (Aliases GPU)) -> TypeM GPU ()
forall rep.
(Checkable rep,
 OpWithAliases (Op rep)
 ~ HostOp (Aliases rep) (SOAC (Aliases rep))) =>
Maybe SegLevel
-> HostOp (Aliases rep) (SOAC (Aliases rep)) -> TypeM rep ()
typeCheckGPUOp Maybe SegLevel
forall a. Maybe a
Nothing
    where
      typeCheckGPUOp :: Maybe SegLevel
-> HostOp (Aliases rep) (SOAC (Aliases rep)) -> TypeM rep ()
typeCheckGPUOp Maybe SegLevel
lvl =
        (SegLevel -> OpWithAliases (Op rep) -> TypeM rep ())
-> Maybe SegLevel
-> (SOAC (Aliases rep) -> TypeM rep ())
-> HostOp (Aliases rep) (SOAC (Aliases rep))
-> TypeM rep ()
forall rep op.
Checkable rep =>
(SegLevel -> OpWithAliases (Op rep) -> TypeM rep ())
-> Maybe SegLevel
-> (op -> TypeM rep ())
-> HostOp (Aliases rep) op
-> TypeM rep ()
typeCheckHostOp (Maybe SegLevel
-> HostOp (Aliases rep) (SOAC (Aliases rep)) -> TypeM rep ()
typeCheckGPUOp (Maybe SegLevel
 -> HostOp (Aliases rep) (SOAC (Aliases rep)) -> TypeM rep ())
-> (SegLevel -> Maybe SegLevel)
-> SegLevel
-> HostOp (Aliases rep) (SOAC (Aliases rep))
-> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegLevel -> Maybe SegLevel
forall a. a -> Maybe a
Just) Maybe SegLevel
lvl SOAC (Aliases rep) -> TypeM rep ()
forall rep. Checkable rep => SOAC (Aliases rep) -> TypeM rep ()
typeCheckSOAC

instance TC.Checkable GPU

instance Buildable GPU where
  mkBody :: Stms GPU -> Result -> Body GPU
mkBody = BodyDec GPU -> Stms GPU -> Result -> Body GPU
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body ()
  mkExpPat :: [Ident] -> Exp GPU -> Pat (LetDec GPU)
mkExpPat [Ident]
idents Exp GPU
_ = [Ident] -> Pat Type
basicPat [Ident]
idents
  mkExpDec :: Pat (LetDec GPU) -> Exp GPU -> ExpDec GPU
mkExpDec Pat (LetDec GPU)
_ Exp GPU
_ = ()
  mkLetNames :: [VName] -> Exp GPU -> m (Stm GPU)
mkLetNames = [VName] -> Exp GPU -> m (Stm GPU)
forall rep (m :: * -> *).
(ExpDec rep ~ (), LetDec rep ~ Type, MonadFreshNames m,
 TypedOp (Op rep), HasScope rep m) =>
[VName] -> Exp rep -> m (Stm rep)
simpleMkLetNames

instance BuilderOps GPU

instance PrettyRep GPU

instance HasSegOp GPU where
  type SegOpLevel GPU = SegLevel
  asSegOp :: Op GPU -> Maybe (SegOp (SegOpLevel GPU) GPU)
asSegOp (SegOp op) = SegOp SegLevel GPU -> Maybe (SegOp SegLevel GPU)
forall a. a -> Maybe a
Just SegOp SegLevel GPU
op
  asSegOp Op GPU
_ = Maybe (SegOp (SegOpLevel GPU) GPU)
forall a. Maybe a
Nothing
  segOp :: SegOp (SegOpLevel GPU) GPU -> Op GPU
segOp = SegOp (SegOpLevel GPU) GPU -> Op GPU
forall rep op. SegOp SegLevel rep -> HostOp rep op
SegOp