{-# LANGUAGE FlexibleContexts #-}

-- | Optimisation pipelines.
module Futhark.Passes
  ( standardPipeline,
    sequentialPipeline,
    kernelsPipeline,
    sequentialCpuPipeline,
    gpuPipeline,
    mcPipeline,
    multicorePipeline,
  )
where

import Control.Category ((>>>))
import Futhark.IR.Kernels (Kernels)
import Futhark.IR.KernelsMem (KernelsMem)
import Futhark.IR.MC (MC)
import Futhark.IR.MCMem (MCMem)
import Futhark.IR.SOACS (SOACS)
import Futhark.IR.Seq (Seq)
import Futhark.IR.SeqMem (SeqMem)
import Futhark.Optimise.CSE
import Futhark.Optimise.DoubleBuffer
import Futhark.Optimise.Fusion
import Futhark.Optimise.InPlaceLowering
import Futhark.Optimise.InliningDeadFun
import Futhark.Optimise.Sink
import Futhark.Optimise.TileLoops
import Futhark.Optimise.Unstream
import Futhark.Pass.ExpandAllocations
import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels
import qualified Futhark.Pass.ExplicitAllocations.MC as MC
import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq
import Futhark.Pass.ExtractKernels
import Futhark.Pass.ExtractMulticore
import Futhark.Pass.FirstOrderTransform
import Futhark.Pass.KernelBabysitting
import Futhark.Pass.Simplify
import Futhark.Pipeline

standardPipeline :: Pipeline SOACS SOACS
standardPipeline :: Pipeline SOACS SOACS
standardPipeline =
  [Pass SOACS SOACS] -> Pipeline SOACS SOACS
forall lore.
Checkable lore =>
[Pass lore lore] -> Pipeline lore lore
passes
    [ Pass SOACS SOACS
simplifySOACS,
      Pass SOACS SOACS
inlineFunctions,
      Pass SOACS SOACS
simplifySOACS,
      Bool -> Pass SOACS SOACS
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
True,
      Pass SOACS SOACS
simplifySOACS,
      -- We run fusion twice
      Pass SOACS SOACS
fuseSOACs,
      Bool -> Pass SOACS SOACS
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
True,
      Pass SOACS SOACS
simplifySOACS,
      Pass SOACS SOACS
fuseSOACs,
      Bool -> Pass SOACS SOACS
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
True,
      Pass SOACS SOACS
simplifySOACS,
      Pass SOACS SOACS
removeDeadFunctions
    ]

kernelsPipeline :: Pipeline SOACS Kernels
kernelsPipeline :: Pipeline SOACS Kernels
kernelsPipeline =
  Pipeline SOACS SOACS
standardPipeline
    Pipeline SOACS SOACS
-> Pipeline SOACS Kernels -> Pipeline SOACS Kernels
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass SOACS Kernels -> Pipeline SOACS Kernels
forall tolore fromlore.
Checkable tolore =>
Pass fromlore tolore -> Pipeline fromlore tolore
onePass Pass SOACS Kernels
extractKernels
    Pipeline SOACS Kernels
-> Pipeline Kernels Kernels -> Pipeline SOACS Kernels
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass Kernels Kernels] -> Pipeline Kernels Kernels
forall lore.
Checkable lore =>
[Pass lore lore] -> Pipeline lore lore
passes
      [ Pass Kernels Kernels
simplifyKernels,
        Pass Kernels Kernels
babysitKernels,
        Pass Kernels Kernels
tileLoops,
        Pass Kernels Kernels
unstreamKernels,
        Bool -> Pass Kernels Kernels
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
True,
        Pass Kernels Kernels
simplifyKernels,
        Pass Kernels Kernels
sinkKernels,
        Pass Kernels Kernels
inPlaceLoweringKernels
      ]

sequentialPipeline :: Pipeline SOACS Seq
sequentialPipeline :: Pipeline SOACS Seq
sequentialPipeline =
  Pipeline SOACS SOACS
standardPipeline
    Pipeline SOACS SOACS -> Pipeline SOACS Seq -> Pipeline SOACS Seq
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass SOACS Seq -> Pipeline SOACS Seq
forall tolore fromlore.
Checkable tolore =>
Pass fromlore tolore -> Pipeline fromlore tolore
onePass Pass SOACS Seq
forall lore. FirstOrderLore lore => Pass SOACS lore
firstOrderTransform
    Pipeline SOACS Seq -> Pipeline Seq Seq -> Pipeline SOACS Seq
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass Seq Seq] -> Pipeline Seq Seq
forall lore.
Checkable lore =>
[Pass lore lore] -> Pipeline lore lore
passes
      [ Pass Seq Seq
simplifySeq,
        Pass Seq Seq
inPlaceLoweringSeq
      ]

sequentialCpuPipeline :: Pipeline SOACS SeqMem
sequentialCpuPipeline :: Pipeline SOACS SeqMem
sequentialCpuPipeline =
  Pipeline SOACS Seq
sequentialPipeline
    Pipeline SOACS Seq -> Pipeline Seq SeqMem -> Pipeline SOACS SeqMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass Seq SeqMem -> Pipeline Seq SeqMem
forall tolore fromlore.
Checkable tolore =>
Pass fromlore tolore -> Pipeline fromlore tolore
onePass Pass Seq SeqMem
Seq.explicitAllocations
    Pipeline Seq SeqMem
-> Pipeline SeqMem SeqMem -> Pipeline Seq SeqMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass SeqMem SeqMem] -> Pipeline SeqMem SeqMem
forall lore.
Checkable lore =>
[Pass lore lore] -> Pipeline lore lore
passes
      [ Bool -> Pass SeqMem SeqMem
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
False,
        Pass SeqMem SeqMem
simplifySeqMem,
        Pass SeqMem SeqMem
simplifySeqMem
      ]

gpuPipeline :: Pipeline SOACS KernelsMem
gpuPipeline :: Pipeline SOACS KernelsMem
gpuPipeline =
  Pipeline SOACS Kernels
kernelsPipeline
    Pipeline SOACS Kernels
-> Pipeline Kernels KernelsMem -> Pipeline SOACS KernelsMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass Kernels KernelsMem -> Pipeline Kernels KernelsMem
forall tolore fromlore.
Checkable tolore =>
Pass fromlore tolore -> Pipeline fromlore tolore
onePass Pass Kernels KernelsMem
Kernels.explicitAllocations
    Pipeline Kernels KernelsMem
-> Pipeline KernelsMem KernelsMem -> Pipeline Kernels KernelsMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass KernelsMem KernelsMem] -> Pipeline KernelsMem KernelsMem
forall lore.
Checkable lore =>
[Pass lore lore] -> Pipeline lore lore
passes
      [ Pass KernelsMem KernelsMem
simplifyKernelsMem,
        Bool -> Pass KernelsMem KernelsMem
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
False,
        Pass KernelsMem KernelsMem
simplifyKernelsMem,
        Pass KernelsMem KernelsMem
doubleBufferKernels,
        Pass KernelsMem KernelsMem
simplifyKernelsMem,
        Pass KernelsMem KernelsMem
expandAllocations,
        Pass KernelsMem KernelsMem
simplifyKernelsMem
      ]

mcPipeline :: Pipeline SOACS MC
mcPipeline :: Pipeline SOACS MC
mcPipeline =
  Pipeline SOACS SOACS
standardPipeline
    Pipeline SOACS SOACS -> Pipeline SOACS MC -> Pipeline SOACS MC
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass SOACS MC -> Pipeline SOACS MC
forall tolore fromlore.
Checkable tolore =>
Pass fromlore tolore -> Pipeline fromlore tolore
onePass Pass SOACS MC
extractMulticore
    Pipeline SOACS MC -> Pipeline MC MC -> Pipeline SOACS MC
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass MC MC] -> Pipeline MC MC
forall lore.
Checkable lore =>
[Pass lore lore] -> Pipeline lore lore
passes
      [ Pass MC MC
simplifyMC,
        Pass MC MC
unstreamMC,
        Bool -> Pass MC MC
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
True,
        Pass MC MC
simplifyMC,
        Pass MC MC
sinkMC,
        Pass MC MC
inPlaceLoweringMC
      ]

multicorePipeline :: Pipeline SOACS MCMem
multicorePipeline :: Pipeline SOACS MCMem
multicorePipeline =
  Pipeline SOACS MC
mcPipeline
    Pipeline SOACS MC -> Pipeline MC MCMem -> Pipeline SOACS MCMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass MC MCMem -> Pipeline MC MCMem
forall tolore fromlore.
Checkable tolore =>
Pass fromlore tolore -> Pipeline fromlore tolore
onePass Pass MC MCMem
MC.explicitAllocations
    Pipeline MC MCMem -> Pipeline MCMem MCMem -> Pipeline MC MCMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass MCMem MCMem] -> Pipeline MCMem MCMem
forall lore.
Checkable lore =>
[Pass lore lore] -> Pipeline lore lore
passes
      [ Pass MCMem MCMem
simplifyMCMem,
        Bool -> Pass MCMem MCMem
forall lore.
(ASTLore lore, CanBeAliased (Op lore),
 CSEInOp (OpWithAliases (Op lore))) =>
Bool -> Pass lore lore
performCSE Bool
False,
        Pass MCMem MCMem
simplifyMCMem,
        Pass MCMem MCMem
doubleBufferMC,
        Pass MCMem MCMem
simplifyMCMem
      ]