{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# 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) where
  opSizeSubst :: PatT dec -> MCOp rep op -> ChunkMap
opSizeSubst PatT dec
_ MCOp rep op
_ = ChunkMap
forall a. Monoid a => a
mempty

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
  SegOpMapper () MC MCMem (AllocM MC MCMem)
-> SegOp () MC -> AllocM MC MCMem (SegOp () MCMem)
forall (m :: * -> *) lvl frep trep.
(Applicative m, 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 = SegSpace -> Scope MCMem
forall rep. SegSpace -> Scope rep
scopeOfSegSpace (SegSpace -> Scope MCMem) -> SegSpace -> Scope MCMem
forall a b. (a -> b) -> a -> b
$ SegOp () MC -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp () MC
op
    mapper :: SubExp -> SegOpMapper () MC MCMem (AllocM MC MCMem)
mapper SubExp
num_threads =
      SegOpMapper () Any Any (AllocM MC MCMem)
forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
identitySegOpMapper
        { mapOnSegOpBody :: KernelBody MC -> AllocM MC MCMem (KernelBody MCMem)
mapOnSegOpBody =
            Scope MCMem
-> AllocM MC MCMem (KernelBody MCMem)
-> AllocM MC MCMem (KernelBody MCMem)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope Scope MCMem
scope (AllocM MC MCMem (KernelBody MCMem)
 -> AllocM MC MCMem (KernelBody MCMem))
-> (KernelBody MC -> AllocM MC MCMem (KernelBody MCMem))
-> KernelBody MC
-> AllocM MC MCMem (KernelBody MCMem)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelBody MC -> AllocM MC MCMem (KernelBody MCMem)
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 =
            SubExp -> SegSpace -> Lambda MC -> AllocM MC MCMem (Lambda MCMem)
forall fromrep torep inner.
Allocable fromrep torep inner =>
SubExp
-> SegSpace
-> Lambda fromrep
-> AllocM fromrep torep (Lambda torep)
allocInBinOpLambda SubExp
num_threads (SegOp () MC -> SegSpace
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 par_op op) =
  MCOp MCMem () -> MemOp (MCOp MCMem ())
forall inner. inner -> MemOp inner
Inner (MCOp MCMem () -> MemOp (MCOp MCMem ()))
-> AllocM MC MCMem (MCOp MCMem ())
-> AllocM MC MCMem (MemOp (MCOp MCMem ()))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Maybe (SegOp () MCMem) -> SegOp () MCMem -> MCOp MCMem ()
forall rep op. Maybe (SegOp () rep) -> SegOp () rep -> MCOp rep op
ParOp (Maybe (SegOp () MCMem) -> SegOp () MCMem -> MCOp MCMem ())
-> AllocM MC MCMem (Maybe (SegOp () MCMem))
-> AllocM MC MCMem (SegOp () MCMem -> MCOp MCMem ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SegOp () MC -> AllocM MC MCMem (SegOp () MCMem))
-> Maybe (SegOp () MC) -> AllocM MC MCMem (Maybe (SegOp () MCMem))
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 AllocM MC MCMem (SegOp () MCMem -> MCOp MCMem ())
-> AllocM MC MCMem (SegOp () MCMem)
-> AllocM MC MCMem (MCOp MCMem ())
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) =
  [Char] -> AllocM MC MCMem (MemOp (MCOp MCMem ()))
forall a. HasCallStack => [Char] -> a
error ([Char] -> AllocM MC MCMem (MemOp (MCOp MCMem ())))
-> [Char] -> AllocM MC MCMem (MemOp (MCOp MCMem ()))
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot allocate memory in SOAC: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SOAC MC -> [Char]
forall a. Pretty a => a -> [Char]
pretty SOAC MC
soac

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