-- | Remove grid information from all SegOp operations - the idea is to remove
-- all GPU-specific information from a GPU IR program.
module Futhark.Pass.NoGrid (noGrid) where

import Control.Monad.Identity
import Futhark.IR.GPU
import Futhark.Pass

noGrid :: Pass GPU GPU
noGrid :: Pass GPU GPU
noGrid =
  String -> String -> (Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU
forall fromrep torep.
String
-> String
-> (Prog fromrep -> PassM (Prog torep))
-> Pass fromrep torep
Pass String
"no-grid" String
"remove grid info from SegOps" ((Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU)
-> (Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU
forall a b. (a -> b) -> a -> b
$
    (Scope GPU -> Stms GPU -> PassM (Stms GPU))
-> Prog GPU -> PassM (Prog GPU)
forall rep.
(Scope rep -> Stms rep -> PassM (Stms rep))
-> Prog rep -> PassM (Prog rep)
intraproceduralTransformation Scope GPU -> Stms GPU -> PassM (Stms GPU)
forall {rep} {op :: * -> *} {f :: * -> *} {f :: * -> *} {p}.
(OpC rep ~ HostOp op, RetType rep ~ ExtType,
 BranchType rep ~ ExtType, LParamInfo rep ~ Type, Applicative f,
 Functor f) =>
p -> f (Stm rep) -> f (f (Stm rep))
optimise
  where
    optimise :: p -> f (Stm rep) -> f (f (Stm rep))
optimise p
_scope f (Stm rep)
stms = f (Stm rep) -> f (f (Stm rep))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (f (Stm rep) -> f (f (Stm rep))) -> f (Stm rep) -> f (f (Stm rep))
forall a b. (a -> b) -> a -> b
$ (Stm rep -> Stm rep) -> f (Stm rep) -> f (Stm rep)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm rep -> Stm rep
forall {rep} {op :: * -> *}.
(RetType rep ~ ExtType, BranchType rep ~ ExtType,
 OpC rep ~ HostOp op, LParamInfo rep ~ Type) =>
Stm rep -> Stm rep
onStm f (Stm rep)
stms
    onStm :: Stm rep -> Stm rep
onStm (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux Exp rep
e) =
      Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$
        Mapper rep rep Identity -> Exp rep -> Exp rep
forall frep trep. Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp
          (forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper @GPU)
            { mapOnOp = pure . onOp,
              mapOnBody = const $ pure . onBody,
              mapOnFParam = pure
            }
          Exp rep
e
    onBody :: Body rep -> Body rep
onBody Body rep
body = Body rep
body {bodyStms = fmap onStm (bodyStms body)}

    onLevel :: SegLevel -> SegLevel
onLevel (SegThread SegVirt
v Maybe KernelGrid
_) = SegVirt -> Maybe KernelGrid -> SegLevel
SegThread SegVirt
v Maybe KernelGrid
forall a. Maybe a
Nothing
    onLevel (SegBlock SegVirt
v Maybe KernelGrid
_) = SegVirt -> Maybe KernelGrid -> SegLevel
SegBlock SegVirt
v Maybe KernelGrid
forall a. Maybe a
Nothing
    onLevel (SegThreadInBlock SegVirt
v) = SegVirt -> SegLevel
SegThreadInBlock SegVirt
v

    onOp :: HostOp op rep -> HostOp op rep
onOp (SegOp SegOp SegLevel rep
op) =
      SegOp SegLevel rep -> HostOp op rep
forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
SegOp (SegOp SegLevel rep -> HostOp op rep)
-> (Identity (SegOp SegLevel rep) -> SegOp SegLevel rep)
-> Identity (SegOp SegLevel rep)
-> HostOp op rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity (SegOp SegLevel rep) -> SegOp SegLevel rep
forall a. Identity a -> a
runIdentity (Identity (SegOp SegLevel rep) -> HostOp op rep)
-> Identity (SegOp SegLevel rep) -> HostOp op rep
forall a b. (a -> b) -> a -> b
$
        SegOpMapper SegLevel rep rep Identity
-> SegOp SegLevel rep -> Identity (SegOp SegLevel rep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM
          (SegOpMapper (ZonkAny 0) rep rep Identity
forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
identitySegOpMapper {mapOnSegOpLevel = pure . onLevel})
          SegOp SegLevel rep
op
    onOp HostOp op rep
op = HostOp op rep
op