    {-# 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 NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid
    KernelGrid
            (Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid)
    -> m (Count NumBlocks SubExp)
    -> m (Count BlockSize SubExp -> KernelGrid)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> Count NumBlocks SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count (SubExp -> Count NumBlocks SubExp)
    -> m SubExp -> m (Count NumBlocks 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_tblocks") SizeClass
    SizeGrid)
            m (Count BlockSize SubExp -> KernelGrid)
    -> m (Count BlockSize 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 BlockSize SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count (SubExp -> Count BlockSize SubExp)
    -> m SubExp -> m (Count BlockSize 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
    "_tblock_size") SizeClass
    SizeThreadBlock)

    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