{-# LANGUAGE TypeFamilies #-} module Futhark.Pass.ExtractKernels.ToGPU ( getSize, segThread, soacsLambdaToGPU, soacsStmToGPU, scopeForGPU, scopeForSOACs, injectSOACS, ) where import Control.Monad.Identity import Data.List () import Futhark.IR import Futhark.IR.GPU import Futhark.IR.SOACS (SOACS) import Futhark.IR.SOACS.SOAC qualified as SOAC import Futhark.Tools getSize :: (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) => String -> SizeClass -> m SubExp getSize :: forall (m :: * -> *) (inner :: * -> *). (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) => String -> SizeClass -> m SubExp getSize String desc SizeClass size_class = do Name size_key <- String -> Name nameFromString (String -> Name) -> (VName -> String) -> VName -> Name forall b c a. (b -> c) -> (a -> b) -> a -> c . VName -> String forall a. Pretty a => a -> String prettyString (VName -> Name) -> m VName -> m Name forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> String -> m VName forall (m :: * -> *). MonadFreshNames m => String -> m VName newVName String desc String -> Exp (Rep m) -> m SubExp forall (m :: * -> *). MonadBuilder m => String -> Exp (Rep m) -> m SubExp letSubExp String desc (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp forall a b. (a -> b) -> a -> b $ Op (Rep m) -> Exp (Rep m) forall rep. Op rep -> Exp rep Op (Op (Rep m) -> Exp (Rep m)) -> Op (Rep m) -> Exp (Rep m) forall a b. (a -> b) -> a -> b $ SizeOp -> HostOp inner (Rep m) forall (op :: * -> *) rep. SizeOp -> HostOp op rep SizeOp (SizeOp -> HostOp inner (Rep m)) -> SizeOp -> HostOp inner (Rep m) forall a b. (a -> b) -> a -> b $ Name -> SizeClass -> SizeOp GetSize Name size_key SizeClass size_class segThread :: (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) => String -> m SegLevel segThread :: forall (m :: * -> *) (inner :: * -> *). (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) => String -> m SegLevel segThread String desc = SegVirt -> Maybe KernelGrid -> SegLevel SegThread SegVirt SegVirt (Maybe KernelGrid -> SegLevel) -> m (Maybe KernelGrid) -> m SegLevel forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> (KernelGrid -> Maybe KernelGrid forall a. a -> Maybe a Just (KernelGrid -> Maybe KernelGrid) -> m KernelGrid -> m (Maybe KernelGrid) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> m KernelGrid kernelGrid) where kernelGrid :: m KernelGrid kernelGrid = Count NumGroups SubExp -> Count GroupSize SubExp -> KernelGrid KernelGrid (Count NumGroups SubExp -> Count GroupSize SubExp -> KernelGrid) -> m (Count NumGroups SubExp) -> m (Count GroupSize SubExp -> KernelGrid) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> (SubExp -> Count NumGroups SubExp forall {k} (u :: k) e. e -> Count u e Count (SubExp -> Count NumGroups SubExp) -> m SubExp -> m (Count NumGroups SubExp) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> String -> SizeClass -> m SubExp forall (m :: * -> *) (inner :: * -> *). (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) => String -> SizeClass -> m SubExp getSize (String desc String -> String -> String forall a. [a] -> [a] -> [a] ++ String "_num_groups") SizeClass SizeNumGroups) m (Count GroupSize SubExp -> KernelGrid) -> m (Count GroupSize SubExp) -> m KernelGrid forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> (SubExp -> Count GroupSize SubExp forall {k} (u :: k) e. e -> Count u e Count (SubExp -> Count GroupSize SubExp) -> m SubExp -> m (Count GroupSize SubExp) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> String -> SizeClass -> m SubExp forall (m :: * -> *) (inner :: * -> *). (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) => String -> SizeClass -> m SubExp getSize (String desc String -> String -> String forall a. [a] -> [a] -> [a] ++ String "_group_size") SizeClass SizeGroup) injectSOACS :: ( Monad m, SameScope from to, ExpDec from ~ ExpDec to, BodyDec from ~ BodyDec to, RetType from ~ RetType to, BranchType from ~ BranchType to, Op from ~ SOAC from ) => (SOAC to -> Op to) -> Rephraser m from to injectSOACS :: forall (m :: * -> *) from to. (Monad m, SameScope from to, ExpDec from ~ ExpDec to, BodyDec from ~ BodyDec to, RetType from ~ RetType to, BranchType from ~ BranchType to, Op from ~ SOAC from) => (SOAC to -> Op to) -> Rephraser m from to injectSOACS SOAC to -> Op to f = Rephraser { rephraseExpDec :: ExpDec from -> m (ExpDec to) rephraseExpDec = ExpDec from -> m (ExpDec to) ExpDec to -> m (ExpDec to) forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, rephraseBodyDec :: BodyDec from -> m (BodyDec to) rephraseBodyDec = BodyDec from -> m (BodyDec to) BodyDec to -> m (BodyDec to) forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, rephraseLetBoundDec :: LetDec from -> m (LetDec to) rephraseLetBoundDec = LetDec from -> m (LetDec to) LetDec to -> m (LetDec to) forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, rephraseFParamDec :: FParamInfo from -> m (FParamInfo to) rephraseFParamDec = FParamInfo from -> m (FParamInfo to) FParamInfo to -> m (FParamInfo to) forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, rephraseLParamDec :: LParamInfo from -> m (LParamInfo to) rephraseLParamDec = LParamInfo from -> m (LParamInfo to) LParamInfo to -> m (LParamInfo to) forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, rephraseOp :: Op from -> m (Op to) rephraseOp = (SOAC to -> Op to) -> m (SOAC to) -> m (Op to) forall a b. (a -> b) -> m a -> m b forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b fmap SOAC to -> Op to f (m (SOAC to) -> m (Op to)) -> (SOAC from -> m (SOAC to)) -> SOAC from -> m (Op to) forall b c a. (b -> c) -> (a -> b) -> a -> c . SOAC from -> m (SOAC to) onSOAC, rephraseRetType :: RetType from -> m (RetType to) rephraseRetType = RetType from -> m (RetType to) RetType to -> m (RetType to) forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, rephraseBranchType :: BranchType from -> m (BranchType to) rephraseBranchType = BranchType from -> m (BranchType to) BranchType to -> m (BranchType to) forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure } where onSOAC :: SOAC from -> m (SOAC to) onSOAC = SOACMapper from to m -> SOAC from -> m (SOAC to) forall (m :: * -> *) frep trep. Monad m => SOACMapper frep trep m -> SOAC frep -> m (SOAC trep) SOAC.mapSOACM SOACMapper from to m mapper mapper :: SOACMapper from to m mapper = SOAC.SOACMapper { mapOnSOACSubExp :: SubExp -> m SubExp SOAC.mapOnSOACSubExp = SubExp -> m SubExp forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, mapOnSOACVName :: VName -> m VName SOAC.mapOnSOACVName = VName -> m VName forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure, mapOnSOACLambda :: Lambda from -> m (Lambda to) SOAC.mapOnSOACLambda = Rephraser m from to -> Lambda from -> m (Lambda to) forall (m :: * -> *) from to. Monad m => Rephraser m from to -> Lambda from -> m (Lambda to) rephraseLambda (Rephraser m from to -> Lambda from -> m (Lambda to)) -> Rephraser m from to -> Lambda from -> m (Lambda to) forall a b. (a -> b) -> a -> b $ (SOAC to -> Op to) -> Rephraser m from to forall (m :: * -> *) from to. (Monad m, SameScope from to, ExpDec from ~ ExpDec to, BodyDec from ~ BodyDec to, RetType from ~ RetType to, BranchType from ~ BranchType to, Op from ~ SOAC from) => (SOAC to -> Op to) -> Rephraser m from to injectSOACS SOAC to -> Op to f } soacsStmToGPU :: Stm SOACS -> Stm GPU soacsStmToGPU :: Stm SOACS -> Stm GPU soacsStmToGPU = Identity (Stm GPU) -> Stm GPU forall a. Identity a -> a runIdentity (Identity (Stm GPU) -> Stm GPU) -> (Stm SOACS -> Identity (Stm GPU)) -> Stm SOACS -> Stm GPU forall b c a. (b -> c) -> (a -> b) -> a -> c . Rephraser Identity SOACS GPU -> Stm SOACS -> Identity (Stm GPU) forall (m :: * -> *) from to. Monad m => Rephraser m from to -> Stm from -> m (Stm to) rephraseStm ((SOAC GPU -> Op GPU) -> Rephraser Identity SOACS GPU forall (m :: * -> *) from to. (Monad m, SameScope from to, ExpDec from ~ ExpDec to, BodyDec from ~ BodyDec to, RetType from ~ RetType to, BranchType from ~ BranchType to, Op from ~ SOAC from) => (SOAC to -> Op to) -> Rephraser m from to injectSOACS SOAC GPU -> Op GPU SOAC GPU -> HostOp SOAC GPU forall (op :: * -> *) rep. op rep -> HostOp op rep OtherOp) soacsLambdaToGPU :: Lambda SOACS -> Lambda GPU soacsLambdaToGPU :: Lambda SOACS -> Lambda GPU soacsLambdaToGPU = Identity (Lambda GPU) -> Lambda GPU forall a. Identity a -> a runIdentity (Identity (Lambda GPU) -> Lambda GPU) -> (Lambda SOACS -> Identity (Lambda GPU)) -> Lambda SOACS -> Lambda GPU forall b c a. (b -> c) -> (a -> b) -> a -> c . Rephraser Identity SOACS GPU -> Lambda SOACS -> Identity (Lambda GPU) forall (m :: * -> *) from to. Monad m => Rephraser m from to -> Lambda from -> m (Lambda to) rephraseLambda ((SOAC GPU -> Op GPU) -> Rephraser Identity SOACS GPU forall (m :: * -> *) from to. (Monad m, SameScope from to, ExpDec from ~ ExpDec to, BodyDec from ~ BodyDec to, RetType from ~ RetType to, BranchType from ~ BranchType to, Op from ~ SOAC from) => (SOAC to -> Op to) -> Rephraser m from to injectSOACS SOAC GPU -> Op GPU SOAC GPU -> HostOp SOAC GPU forall (op :: * -> *) rep. op rep -> HostOp op rep OtherOp) scopeForSOACs :: Scope GPU -> Scope SOACS scopeForSOACs :: Scope GPU -> Scope SOACS scopeForSOACs = Scope GPU -> Scope SOACS forall fromrep torep. SameScope fromrep torep => Scope fromrep -> Scope torep castScope scopeForGPU :: Scope SOACS -> Scope GPU scopeForGPU :: Scope SOACS -> Scope GPU scopeForGPU = Scope SOACS -> Scope GPU forall fromrep torep. SameScope fromrep torep => Scope fromrep -> Scope torep castScope