{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

-- | Converting 'MC' programs to 'MCMem'.
module Futhark.Pass.ExplicitAllocations.MC (explicitAllocations) where

import Futhark.IR.MC
import Futhark.IR.MCMem
import Futhark.Pass.ExplicitAllocations
import Futhark.Pass.ExplicitAllocations.SegOp

instance SizeSubst (MCOp rep op)

handleSegOp :: SegOp () MC -> AllocM MC MCMem (SegOp () MCMem)
handleSegOp :: SegOp () MC -> AllocM MC MCMem (SegOp () MCMem)
handleSegOp SegOp () MC
op = do
  let num_threads :: SubExp
num_threads = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
256 -- FIXME
  forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM (SubExp -> SegOpMapper () MC MCMem (AllocM MC MCMem)
mapper SubExp
num_threads) SegOp () MC
op
  where
    scope :: Scope MCMem
scope = forall rep. SegSpace -> Scope rep
scopeOfSegSpace forall a b. (a -> b) -> a -> b
$ forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp () MC
op
    mapper :: SubExp -> SegOpMapper () MC MCMem (AllocM MC MCMem)
mapper SubExp
num_threads =
      forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
identitySegOpMapper
        { mapOnSegOpBody :: KernelBody MC -> AllocM MC MCMem (KernelBody MCMem)
mapOnSegOpBody =
            forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope Scope MCMem
scope forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fromrep torep (inner :: * -> *).
Allocable fromrep torep inner =>
KernelBody fromrep -> AllocM fromrep torep (KernelBody torep)
allocInKernelBody,
          mapOnSegOpLambda :: Lambda MC -> AllocM MC MCMem (Lambda MCMem)
mapOnSegOpLambda =
            forall fromrep torep (inner :: * -> *).
Allocable fromrep torep inner =>
SubExp
-> SegSpace
-> Lambda fromrep
-> AllocM fromrep torep (Lambda torep)
allocInBinOpLambda SubExp
num_threads (forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp () MC
op)
        }

handleMCOp :: Op MC -> AllocM MC MCMem (Op MCMem)
handleMCOp :: Op MC -> AllocM MC MCMem (Op MCMem)
handleMCOp (ParOp Maybe (SegOp () MC)
par_op SegOp () MC
op) =
  forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Inner forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse SegOp () MC -> AllocM MC MCMem (SegOp () MCMem)
handleSegOp Maybe (SegOp () MC)
par_op forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOp () MC -> AllocM MC MCMem (SegOp () MCMem)
handleSegOp SegOp () MC
op)
handleMCOp (OtherOp SOAC MC
soac) =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot allocate memory in SOAC: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString SOAC MC
soac

-- | The pass from 'MC' to 'MCMem'.
explicitAllocations :: Pass MC MCMem
explicitAllocations :: Pass MC MCMem
explicitAllocations = forall fromrep torep (inner :: * -> *).
Allocable fromrep torep inner =>
Space
-> (Op fromrep -> AllocM fromrep torep (Op torep))
-> (Exp torep -> AllocM fromrep torep [ExpHint])
-> Pass fromrep torep
explicitAllocationsGeneric Space
DefaultSpace Op MC -> AllocM MC MCMem (Op MCMem)
handleMCOp forall rep (m :: * -> *).
(ASTRep rep, HasScope rep m) =>
Exp rep -> m [ExpHint]
defaultExpHints