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

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 lore op) where
  opSizeSubst :: PatternT dec -> MCOp lore op -> ChunkMap
opSizeSubst PatternT dec
_ MCOp lore 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 flore tlore.
(Applicative m, Monad m) =>
SegOpMapper lvl flore tlore m
-> SegOp lvl flore -> m (SegOp lvl tlore)
mapSegOpM (SubExp -> SegOpMapper () MC MCMem (AllocM MC MCMem)
mapper SubExp
num_threads) SegOp () MC
op
  where
    scope :: Scope MCMem
scope = SegSpace -> Scope MCMem
forall lore. SegSpace -> Scope lore
scopeOfSegSpace (SegSpace -> Scope MCMem) -> SegSpace -> Scope MCMem
forall a b. (a -> b) -> a -> b
$ SegOp () MC -> SegSpace
forall lvl lore. SegOp lvl lore -> 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 lore.
Monad m =>
SegOpMapper lvl lore lore m
identitySegOpMapper
        { mapOnSegOpBody :: KernelBody MC -> AllocM MC MCMem (KernelBody MCMem)
mapOnSegOpBody =
            Scope MCMem
-> AllocM MC MCMem (KernelBody MCMem)
-> AllocM MC MCMem (KernelBody MCMem)
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> 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 fromlore tolore.
Allocable fromlore tolore =>
KernelBody fromlore -> AllocM fromlore tolore (KernelBody tolore)
allocInKernelBody,
          mapOnSegOpLambda :: Lambda MC -> AllocM MC MCMem (Lambda MCMem)
mapOnSegOpLambda =
            SubExp -> SegSpace -> Lambda MC -> AllocM MC MCMem (Lambda MCMem)
forall fromlore tolore.
Allocable fromlore tolore =>
SubExp
-> SegSpace
-> Lambda fromlore
-> AllocM fromlore tolore (Lambda tolore)
allocInBinOpLambda SubExp
num_threads (SegOp () MC -> SegSpace
forall lvl lore. SegOp lvl lore -> 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 lore op.
Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore 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

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 fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
(Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> Pass fromlore tolore
explicitAllocationsGeneric Op MC -> AllocM MC MCMem (Op MCMem)
handleMCOp Exp MCMem -> AllocM MC MCMem [ExpHint]
forall (m :: * -> *) lore.
(Monad m, ASTLore lore) =>
Exp lore -> m [ExpHint]
defaultExpHints