{-# LANGUAGE FlexibleContexts #-}

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

import Control.Category ((>>>))
import Futhark.IR.GPU (GPU)
import Futhark.IR.GPUMem (GPUMem)
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 qualified Futhark.Optimise.MemoryBlockMerging as MemoryBlockMerging
import Futhark.Optimise.Sink
import Futhark.Optimise.TileLoops
import Futhark.Optimise.Unstream
import Futhark.Pass.ExpandAllocations
import qualified Futhark.Pass.ExplicitAllocations.GPU as GPU
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

-- | A pipeline used by all current compilers.  Performs inlining,
-- fusion, and various forms of cleanup.  This pipeline will be
-- followed by another one that deals with parallelism and memory.
standardPipeline :: Pipeline SOACS SOACS
standardPipeline :: Pipeline SOACS SOACS
standardPipeline =
  [Pass SOACS SOACS] -> Pipeline SOACS SOACS
forall rep. Checkable rep => [Pass rep rep] -> Pipeline rep rep
passes
    [ Pass SOACS SOACS
simplifySOACS,
      Pass SOACS SOACS
inlineConservatively,
      Pass SOACS SOACS
simplifySOACS,
      Pass SOACS SOACS
inlineAggressively,
      Pass SOACS SOACS
simplifySOACS,
      Bool -> Pass SOACS SOACS
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
True,
      Pass SOACS SOACS
simplifySOACS,
      -- We run fusion twice
      Pass SOACS SOACS
fuseSOACs,
      Bool -> Pass SOACS SOACS
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
True,
      Pass SOACS SOACS
simplifySOACS,
      Pass SOACS SOACS
fuseSOACs,
      Bool -> Pass SOACS SOACS
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
True,
      Pass SOACS SOACS
simplifySOACS,
      Pass SOACS SOACS
removeDeadFunctions
    ]

-- | The pipeline used by the CUDA and OpenCL backends, but before
-- adding memory information.  Includes 'standardPipeline'.
kernelsPipeline :: Pipeline SOACS GPU
kernelsPipeline :: Pipeline SOACS GPU
kernelsPipeline =
  Pipeline SOACS SOACS
standardPipeline
    Pipeline SOACS SOACS -> Pipeline SOACS GPU -> Pipeline SOACS GPU
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass SOACS GPU -> Pipeline SOACS GPU
forall torep fromrep.
Checkable torep =>
Pass fromrep torep -> Pipeline fromrep torep
onePass Pass SOACS GPU
extractKernels
    Pipeline SOACS GPU -> Pipeline GPU GPU -> Pipeline SOACS GPU
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass GPU GPU] -> Pipeline GPU GPU
forall rep. Checkable rep => [Pass rep rep] -> Pipeline rep rep
passes
      [ Pass GPU GPU
simplifyGPU,
        Pass GPU GPU
babysitKernels,
        Pass GPU GPU
tileLoops,
        Pass GPU GPU
unstreamGPU,
        Bool -> Pass GPU GPU
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
True,
        Pass GPU GPU
simplifyGPU,
        Pass GPU GPU
sinkGPU,
        Pass GPU GPU
inPlaceLoweringGPU
      ]

-- | The pipeline used by the sequential backends.  Turns all
-- parallelism into sequential loops.  Includes 'standardPipeline'.
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 torep fromrep.
Checkable torep =>
Pass fromrep torep -> Pipeline fromrep torep
onePass Pass SOACS Seq
forall rep. FirstOrderRep rep => Pass SOACS rep
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 rep. Checkable rep => [Pass rep rep] -> Pipeline rep rep
passes
      [ Pass Seq Seq
simplifySeq,
        Pass Seq Seq
inPlaceLoweringSeq
      ]

-- | Run 'sequentialPipeline', then add memory information (and
-- optimise it slightly).
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 torep fromrep.
Checkable torep =>
Pass fromrep torep -> Pipeline fromrep torep
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 rep. Checkable rep => [Pass rep rep] -> Pipeline rep rep
passes
      [ Bool -> Pass SeqMem SeqMem
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
False,
        Pass SeqMem SeqMem
simplifySeqMem
      ]

-- | Run 'kernelsPipeline', then add memory information (and optimise
-- it a lot).
gpuPipeline :: Pipeline SOACS GPUMem
gpuPipeline :: Pipeline SOACS GPUMem
gpuPipeline =
  Pipeline SOACS GPU
kernelsPipeline
    Pipeline SOACS GPU -> Pipeline GPU GPUMem -> Pipeline SOACS GPUMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Pass GPU GPUMem -> Pipeline GPU GPUMem
forall torep fromrep.
Checkable torep =>
Pass fromrep torep -> Pipeline fromrep torep
onePass Pass GPU GPUMem
GPU.explicitAllocations
    Pipeline GPU GPUMem
-> Pipeline GPUMem GPUMem -> Pipeline GPU GPUMem
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> [Pass GPUMem GPUMem] -> Pipeline GPUMem GPUMem
forall rep. Checkable rep => [Pass rep rep] -> Pipeline rep rep
passes
      [ Pass GPUMem GPUMem
simplifyGPUMem,
        Bool -> Pass GPUMem GPUMem
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
False,
        Pass GPUMem GPUMem
simplifyGPUMem,
        Pass GPUMem GPUMem
doubleBufferGPU,
        Pass GPUMem GPUMem
simplifyGPUMem,
        Pass GPUMem GPUMem
MemoryBlockMerging.optimise,
        Pass GPUMem GPUMem
simplifyGPUMem,
        Pass GPUMem GPUMem
expandAllocations,
        Pass GPUMem GPUMem
simplifyGPUMem
      ]

-- | Run 'standardPipeline' and then convert to multicore
-- representation (and do a bunch of optimisation).
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 torep fromrep.
Checkable torep =>
Pass fromrep torep -> Pipeline fromrep torep
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 rep. Checkable rep => [Pass rep rep] -> Pipeline rep rep
passes
      [ Pass MC MC
simplifyMC,
        Pass MC MC
unstreamMC,
        Bool -> Pass MC MC
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
True,
        Pass MC MC
simplifyMC,
        Pass MC MC
sinkMC,
        Pass MC MC
inPlaceLoweringMC
      ]

-- | Run 'mcPipeline' and then add memory information.
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 torep fromrep.
Checkable torep =>
Pass fromrep torep -> Pipeline fromrep torep
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 rep. Checkable rep => [Pass rep rep] -> Pipeline rep rep
passes
      [ Pass MCMem MCMem
simplifyMCMem,
        Bool -> Pass MCMem MCMem
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Pass rep rep
performCSE Bool
False,
        Pass MCMem MCMem
simplifyMCMem,
        Pass MCMem MCMem
doubleBufferMC,
        Pass MCMem MCMem
simplifyMCMem
      ]