content: module Futhark.Optimise.ArrayLayout
    ( optimiseArrayLayoutGPU,
      optimiseArrayLayoutMC,
    )
  where
  
  import Control.Monad.State.Strict
  import Futhark.Analysis.AccessPattern (Analyse, analyseDimAccesses)
  import Futhark.Analysis.PrimExp.Table (primExpTable)
  import Futhark.Builder
  import Futhark.IR.GPU (GPU)
  import Futhark.IR.MC (MC)
  import Futhark.Optimise.ArrayLayout.Layout (layoutTableFromIndexTable)
  import Futhark.Optimise.ArrayLayout.Transform (Transform, transformStms)
  import Futhark.Pass
  
  optimiseArrayLayout :: (Analyse rep, Transform rep, BuilderOps rep) => String -> Pass rep rep
  optimiseArrayLayout s =
    Pass
      ("optimise array layout " <> s)
      "Transform array layout for locality optimisations."
      $ \prog -> do
        -- Analyse the program
        let index_table = analyseDimAccesses prog
        -- Compute primExps for all variables
        let table = primExpTable prog
        -- Compute permutations to acheive coalescence for all arrays
        let permutation_table = layoutTableFromIndexTable table index_table
        -- Insert permutations in the AST
        intraproceduralTransformation (onStms permutation_table) prog
    where
      onStms layout_table scope stms = do
        let m = transformStms layout_table mempty stms
        fmap fst $ modifyNameSource $ runState $ runBuilderT m scope
  
  -- | The optimisation performed on the GPU representation.
  optimiseArrayLayoutGPU :: Pass GPU GPU
  optimiseArrayLayoutGPU = optimiseArrayLayout "gpu"
  
  -- | The optimisation performed on the MC representation.
  optimiseArrayLayoutMC :: Pass MC MC
  optimiseArrayLayoutMC = optimiseArrayLayout "mc"
  