{- |
Module      :  Camfort.Specification.Stencils.Generate
Description :  Generate stencils for inference and synthesis
Copyright   :  (c) 2017, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish
License     :  Apache-2.0

Maintainer  :  dom.orchard@gmail.com
Stability   :  experimental
-}

{-# LANGUAGE ConstraintKinds           #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE PatternGuards             #-}
{-# LANGUAGE ScopedTypeVariables       #-}

module Camfort.Specification.Stencils.Generate
  (
    EvalLog
  , Neighbour(..)
  , extractRelevantIVS
  , genOffsets
  , genSpecifications
  , isArraySubscript
  , neighbourIndex
  , runStencilInferer
   -- Various helpers that get used by other tools, e.g., array-analysis
  , isVariableExpr
  , convIxToNeighbour
  , indicesToRelativisedOffsets
  , indicesToSpec
  , neighbourToOffset
  , relativise
  , consistentIVSuse
  ) where

import           Control.Monad (void, when, zipWithM)
import           Control.Monad.State.Strict (State, get, put, runState)
import           Control.Monad.Writer.Lazy (WriterT, runWriterT, tell)
import           Control.Monad.Reader (ReaderT, runReaderT, asks)
import           Data.Data (Data)
import           Data.Foldable (foldrM)
import           Data.Generics.Uniplate.Operations (transformBi, universeBi)
import           Data.Graph.Inductive.Graph (lab, pre)
import qualified Data.IntMap as IM
import qualified Data.Map as M
import           Data.Maybe (fromJust, fromMaybe, isJust, mapMaybe)
import           Data.Monoid ((<>))
import qualified Data.Set as S

import qualified Language.Fortran.AST               as F
import qualified Language.Fortran.Analysis          as FA
import qualified Language.Fortran.Analysis.DataFlow as FAD
import qualified Language.Fortran.Util.Position     as FU

import           Camfort.Helpers (collect)
import qualified Camfort.Helpers.Vec as V
import           Camfort.Specification.Stencils.Annotation ()
import           Camfort.Specification.Stencils.Analysis
import           Camfort.Specification.Stencils.InferenceBackend
import           Camfort.Specification.Stencils.Model
  (Approximation(..), Multiplicity(..))
import           Camfort.Specification.Stencils.Syntax
  ( absoluteRep
  , fromBool
  , groupKeyBy
  , hasDuplicates
  , isEmpty
  , isUnit
  , setLinearity
  , Specification(..)
  , Variable)

type Indices a = [[F.Index (FA.Analysis a)]]

type EvalLog = [(String, Variable)]

data SIEnv ann = SIEnv
  {
    -- | In-scope induction variables.
    SIEnv ann -> [Variable]
sieIvs :: [Variable]
  , SIEnv ann -> FlowsGraph ann
sieFlowsGraph :: FAD.FlowsGraph ann
  }

-- | Analysis for working with low-level stencil inference.
type StencilInferer ann = ReaderT (SIEnv ann) (WriterT EvalLog StencilsAnalysis)

-- | Get the list of in-scope induction variables.
getIvs :: StencilInferer ann [Variable]
getIvs :: StencilInferer ann [Variable]
getIvs = (SIEnv ann -> [Variable]) -> StencilInferer ann [Variable]
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SIEnv ann -> [Variable]
forall ann. SIEnv ann -> [Variable]
sieIvs

-- | Get the FlowsGraph for the current analysis.
getFlowsGraph :: StencilInferer ann (FAD.FlowsGraph ann)
getFlowsGraph :: StencilInferer ann (FlowsGraph ann)
getFlowsGraph = (SIEnv ann -> FlowsGraph ann)
-> StencilInferer ann (FlowsGraph ann)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SIEnv ann -> FlowsGraph ann
forall ann. SIEnv ann -> FlowsGraph ann
sieFlowsGraph

runStencilInferer :: StencilInferer ann a -> [Variable] -> FAD.FlowsGraph ann -> StencilsAnalysis (a, EvalLog)
runStencilInferer :: StencilInferer ann a
-> [Variable] -> FlowsGraph ann -> StencilsAnalysis (a, EvalLog)
runStencilInferer StencilInferer ann a
si [Variable]
ivs FlowsGraph ann
flowsGraph = do
  let senv :: SIEnv ann
senv = SIEnv :: forall ann. [Variable] -> FlowsGraph ann -> SIEnv ann
SIEnv { sieIvs :: [Variable]
sieIvs = [Variable]
ivs, sieFlowsGraph :: FlowsGraph ann
sieFlowsGraph = FlowsGraph ann
flowsGraph }
  WriterT EvalLog StencilsAnalysis a -> StencilsAnalysis (a, EvalLog)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT EvalLog StencilsAnalysis a
 -> StencilsAnalysis (a, EvalLog))
-> WriterT EvalLog StencilsAnalysis a
-> StencilsAnalysis (a, EvalLog)
forall a b. (a -> b) -> a -> b
$ StencilInferer ann a
-> SIEnv ann -> WriterT EvalLog StencilsAnalysis a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT StencilInferer ann a
si SIEnv ann
senv

{-| Representation for indices as either:
     * neighbour indices
     * constant
     * non neighbour index -}
data Neighbour = Neighbour Variable Int
               | Constant (F.Value ())
               | NonNeighbour deriving (Neighbour -> Neighbour -> Bool
(Neighbour -> Neighbour -> Bool)
-> (Neighbour -> Neighbour -> Bool) -> Eq Neighbour
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Neighbour -> Neighbour -> Bool
$c/= :: Neighbour -> Neighbour -> Bool
== :: Neighbour -> Neighbour -> Bool
$c== :: Neighbour -> Neighbour -> Bool
Eq, Int -> Neighbour -> ShowS
[Neighbour] -> ShowS
Neighbour -> Variable
(Int -> Neighbour -> ShowS)
-> (Neighbour -> Variable)
-> ([Neighbour] -> ShowS)
-> Show Neighbour
forall a.
(Int -> a -> ShowS) -> (a -> Variable) -> ([a] -> ShowS) -> Show a
showList :: [Neighbour] -> ShowS
$cshowList :: [Neighbour] -> ShowS
show :: Neighbour -> Variable
$cshow :: Neighbour -> Variable
showsPrec :: Int -> Neighbour -> ShowS
$cshowsPrec :: Int -> Neighbour -> ShowS
Show)


{-| Match expressions which are array subscripts, returning Just of their
    index expressions, else Nothing -}
isArraySubscript :: F.Expression (FA.Analysis a) -> Maybe [F.Index (FA.Analysis a)]
isArraySubscript :: Expression (Analysis a) -> Maybe [Index (Analysis a)]
isArraySubscript (F.ExpSubscript Analysis a
_ SrcSpan
_ (F.ExpValue Analysis a
_ SrcSpan
_ (F.ValVariable Variable
_)) AList Index (Analysis a)
subs) =
   [Index (Analysis a)] -> Maybe [Index (Analysis a)]
forall a. a -> Maybe a
Just ([Index (Analysis a)] -> Maybe [Index (Analysis a)])
-> [Index (Analysis a)] -> Maybe [Index (Analysis a)]
forall a b. (a -> b) -> a -> b
$ AList Index (Analysis a) -> [Index (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
F.aStrip AList Index (Analysis a)
subs
isArraySubscript (F.ExpDataRef Analysis a
_ SrcSpan
_ Expression (Analysis a)
e Expression (Analysis a)
e') =
   Expression (Analysis a) -> Maybe [Index (Analysis a)]
forall a. Expression (Analysis a) -> Maybe [Index (Analysis a)]
isArraySubscript Expression (Analysis a)
e Maybe [Index (Analysis a)]
-> Maybe [Index (Analysis a)] -> Maybe [Index (Analysis a)]
forall a. Semigroup a => a -> a -> a
<> Expression (Analysis a) -> Maybe [Index (Analysis a)]
forall a. Expression (Analysis a) -> Maybe [Index (Analysis a)]
isArraySubscript Expression (Analysis a)
e'
isArraySubscript Expression (Analysis a)
_ = Maybe [Index (Analysis a)]
forall a. Maybe a
Nothing

{-| Given an induction-variable-map, convert a list of indices to
    Maybe a list of constant or neighbourhood indices.
    If any are non neighbourhood then return Nothing -}
neighbourIndex :: (Data a)
               => FAD.InductionVarMapByASTBlock
               -> [F.Index (FA.Analysis a)]
               -> Maybe [Neighbour]
neighbourIndex :: InductionVarMapByASTBlock
-> [Index (Analysis a)] -> Maybe [Neighbour]
neighbourIndex InductionVarMapByASTBlock
ivs [Index (Analysis a)]
ixs =
  if Neighbour
NonNeighbour Neighbour -> [Neighbour] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Neighbour]
neighbours
  then [Neighbour] -> Maybe [Neighbour]
forall a. a -> Maybe a
Just [Neighbour]
neighbours
  else Maybe [Neighbour]
forall a. Maybe a
Nothing
    where
      neighbours :: [Neighbour]
neighbours = (Index (Analysis a) -> Neighbour)
-> [Index (Analysis a)] -> [Neighbour]
forall a b. (a -> b) -> [a] -> [b]
map (\Index (Analysis a)
ix -> [Variable] -> Index (Analysis a) -> Neighbour
forall a. Data a => [Variable] -> Index (Analysis a) -> Neighbour
convIxToNeighbour (InductionVarMapByASTBlock -> Index (Analysis a) -> [Variable]
forall (ast :: * -> *) a.
(Spanned (ast (Analysis a)), Annotated ast) =>
InductionVarMapByASTBlock -> ast (Analysis a) -> [Variable]
extractRelevantIVS InductionVarMapByASTBlock
ivs Index (Analysis a)
ix) Index (Analysis a)
ix) [Index (Analysis a)]
ixs

genSpecifications
  :: (Data a, Show a, Eq a)
  => [Neighbour]
  -> F.Block (FA.Analysis a)
  -> StencilInferer a ([([Variable], Specification)], [Int])
genSpecifications :: [Neighbour]
-> Block (Analysis a)
-> StencilInferer a ([([Variable], Specification)], [Int])
genSpecifications [Neighbour]
lhs Block (Analysis a)
block = do
  (Map Variable (Indices a)
subscripts, [Int]
visitedNodes) <- [Block (Analysis a)]
-> StencilInferer a (Map Variable (Indices a), [Int])
forall a.
(Data a, Show a, Eq a) =>
[Block (Analysis a)]
-> StencilInferer a (Map Variable (Indices a), [Int])
genSubscripts [Block (Analysis a)
block]
  [(Variable, Specification)]
varToSpecs <- Map
  Variable
  (ReaderT
     (SIEnv a) (WriterT EvalLog StencilsAnalysis) (Maybe Specification))
-> ReaderT
     (SIEnv a)
     (WriterT EvalLog StencilsAnalysis)
     [(Variable, Specification)]
forall (m :: * -> *) k a.
Monad m =>
Map k (m (Maybe a)) -> m [(k, a)]
assocsSequence (Map
   Variable
   (ReaderT
      (SIEnv a) (WriterT EvalLog StencilsAnalysis) (Maybe Specification))
 -> ReaderT
      (SIEnv a)
      (WriterT EvalLog StencilsAnalysis)
      [(Variable, Specification)])
-> (Map Variable (Indices a)
    -> Map
         Variable
         (ReaderT
            (SIEnv a)
            (WriterT EvalLog StencilsAnalysis)
            (Maybe Specification)))
-> Map Variable (Indices a)
-> ReaderT
     (SIEnv a)
     (WriterT EvalLog StencilsAnalysis)
     [(Variable, Specification)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Variable (Indices a)
-> Map
     Variable
     (ReaderT
        (SIEnv a) (WriterT EvalLog StencilsAnalysis) (Maybe Specification))
mkSpecs (Map Variable (Indices a)
 -> ReaderT
      (SIEnv a)
      (WriterT EvalLog StencilsAnalysis)
      [(Variable, Specification)])
-> Map Variable (Indices a)
-> ReaderT
     (SIEnv a)
     (WriterT EvalLog StencilsAnalysis)
     [(Variable, Specification)]
forall a b. (a -> b) -> a -> b
$ Map Variable (Indices a)
subscripts
  case [(Variable, Specification)]
varToSpecs of
    [] -> do
       EvalLog -> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [(Variable
"EVALMODE: Empty specification (tag: emptySpec)", Variable
"")]
       ([([Variable], Specification)], [Int])
-> StencilInferer a ([([Variable], Specification)], [Int])
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [Int]
visitedNodes)
    [(Variable, Specification)]
_ -> do
       let varsToSpecs :: [([Variable], Specification)]
varsToSpecs = [(Variable, Specification)] -> [([Variable], Specification)]
forall b a. Eq b => [(a, b)] -> [([a], b)]
groupKeyBy [(Variable, Specification)]
varToSpecs
       ([([Variable], Specification)], [Int])
-> StencilInferer a ([([Variable], Specification)], [Int])
forall (m :: * -> *) a. Monad m => a -> m a
return ([([Variable], Specification)] -> [([Variable], Specification)]
forall a. [(a, Specification)] -> [(a, Specification)]
splitUpperAndLower [([Variable], Specification)]
varsToSpecs, [Int]
visitedNodes)
  where
    mkSpecs :: Map Variable (Indices a)
-> Map
     Variable
     (ReaderT
        (SIEnv a) (WriterT EvalLog StencilsAnalysis) (Maybe Specification))
mkSpecs = (Variable
 -> Indices a
 -> ReaderT
      (SIEnv a) (WriterT EvalLog StencilsAnalysis) (Maybe Specification))
-> Map Variable (Indices a)
-> Map
     Variable
     (ReaderT
        (SIEnv a) (WriterT EvalLog StencilsAnalysis) (Maybe Specification))
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (Variable
-> [Neighbour]
-> Indices a
-> ReaderT
     (SIEnv a) (WriterT EvalLog StencilsAnalysis) (Maybe Specification)
forall a.
Data a =>
Variable
-> [Neighbour]
-> Indices a
-> StencilInferer a (Maybe Specification)
`indicesToSpec` [Neighbour]
lhs)

    splitUpperAndLower :: [(a, Specification)] -> [(a, Specification)]
splitUpperAndLower = ((a, Specification) -> [(a, Specification)])
-> [(a, Specification)] -> [(a, Specification)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (a, Specification) -> [(a, Specification)]
forall a. (a, Specification) -> [(a, Specification)]
splitUpperAndLower'
    splitUpperAndLower' :: (a, Specification) -> [(a, Specification)]
splitUpperAndLower' (a
vs, Specification (Mult (Bound (Just Spatial
l) (Just Spatial
u))) Bool
isStencil)
      | Spatial -> Bool
forall t. RegionRig t => t -> Bool
isUnit Spatial
l =
       [(a
vs, Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification (Approximation Spatial -> Multiplicity (Approximation Spatial)
forall a. a -> Multiplicity a
Mult (Maybe Spatial -> Maybe Spatial -> Approximation Spatial
forall a. Maybe a -> Maybe a -> Approximation a
Bound Maybe Spatial
forall a. Maybe a
Nothing (Spatial -> Maybe Spatial
forall a. a -> Maybe a
Just Spatial
u))) Bool
isStencil)]
      | Bool
otherwise =
       [(a
vs, Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification (Approximation Spatial -> Multiplicity (Approximation Spatial)
forall a. a -> Multiplicity a
Mult (Maybe Spatial -> Maybe Spatial -> Approximation Spatial
forall a. Maybe a -> Maybe a -> Approximation a
Bound (Spatial -> Maybe Spatial
forall a. a -> Maybe a
Just Spatial
l) Maybe Spatial
forall a. Maybe a
Nothing)) Bool
isStencil),
        (a
vs, Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification (Approximation Spatial -> Multiplicity (Approximation Spatial)
forall a. a -> Multiplicity a
Mult (Maybe Spatial -> Maybe Spatial -> Approximation Spatial
forall a. Maybe a -> Maybe a -> Approximation a
Bound Maybe Spatial
forall a. Maybe a
Nothing (Spatial -> Maybe Spatial
forall a. a -> Maybe a
Just Spatial
u))) Bool
isStencil)]
    splitUpperAndLower' (a
vs, Specification (Once (Bound (Just Spatial
l) (Just Spatial
u))) Bool
isStencil)
      | Spatial -> Bool
forall t. RegionRig t => t -> Bool
isUnit Spatial
l =
       [(a
vs, Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification (Approximation Spatial -> Multiplicity (Approximation Spatial)
forall a. a -> Multiplicity a
Mult (Maybe Spatial -> Maybe Spatial -> Approximation Spatial
forall a. Maybe a -> Maybe a -> Approximation a
Bound Maybe Spatial
forall a. Maybe a
Nothing (Spatial -> Maybe Spatial
forall a. a -> Maybe a
Just Spatial
u))) Bool
isStencil)]
      | Bool
otherwise =
       [(a
vs, Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification (Approximation Spatial -> Multiplicity (Approximation Spatial)
forall a. a -> Multiplicity a
Once (Maybe Spatial -> Maybe Spatial -> Approximation Spatial
forall a. Maybe a -> Maybe a -> Approximation a
Bound (Spatial -> Maybe Spatial
forall a. a -> Maybe a
Just Spatial
l) Maybe Spatial
forall a. Maybe a
Nothing)) Bool
isStencil),
        (a
vs, Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification (Approximation Spatial -> Multiplicity (Approximation Spatial)
forall a. a -> Multiplicity a
Once (Maybe Spatial -> Maybe Spatial -> Approximation Spatial
forall a. Maybe a -> Maybe a -> Approximation a
Bound Maybe Spatial
forall a. Maybe a
Nothing (Spatial -> Maybe Spatial
forall a. a -> Maybe a
Just Spatial
u))) Bool
isStencil)]
    splitUpperAndLower' (a, Specification)
x = [(a, Specification)
x]

genOffsets
  :: (Data a, Show a, Eq a)
  => [Neighbour]
  -> [F.Block (FA.Analysis a)]
  -> StencilInferer a [(Variable, (Bool, [[Int]]))]
genOffsets :: [Neighbour]
-> [Block (Analysis a)]
-> StencilInferer a [(Variable, (Bool, [[Int]]))]
genOffsets [Neighbour]
lhs [Block (Analysis a)]
blocks = do
  (Map Variable (Indices a)
subscripts, [Int]
_) <- [Block (Analysis a)]
-> StencilInferer a (Map Variable (Indices a), [Int])
forall a.
(Data a, Show a, Eq a) =>
[Block (Analysis a)]
-> StencilInferer a (Map Variable (Indices a), [Int])
genSubscripts [Block (Analysis a)]
blocks
  Map
  Variable
  (ReaderT
     (SIEnv a)
     (WriterT EvalLog StencilsAnalysis)
     (Maybe (Bool, [[Int]])))
-> StencilInferer a [(Variable, (Bool, [[Int]]))]
forall (m :: * -> *) k a.
Monad m =>
Map k (m (Maybe a)) -> m [(k, a)]
assocsSequence (Map
   Variable
   (ReaderT
      (SIEnv a)
      (WriterT EvalLog StencilsAnalysis)
      (Maybe (Bool, [[Int]])))
 -> StencilInferer a [(Variable, (Bool, [[Int]]))])
-> (Map Variable (Indices a)
    -> Map
         Variable
         (ReaderT
            (SIEnv a)
            (WriterT EvalLog StencilsAnalysis)
            (Maybe (Bool, [[Int]]))))
-> Map Variable (Indices a)
-> StencilInferer a [(Variable, (Bool, [[Int]]))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Variable (Indices a)
-> Map
     Variable
     (ReaderT
        (SIEnv a)
        (WriterT EvalLog StencilsAnalysis)
        (Maybe (Bool, [[Int]])))
mkOffsets (Map Variable (Indices a)
 -> StencilInferer a [(Variable, (Bool, [[Int]]))])
-> Map Variable (Indices a)
-> StencilInferer a [(Variable, (Bool, [[Int]]))]
forall a b. (a -> b) -> a -> b
$ Map Variable (Indices a)
subscripts
  where
    mkOffsets :: Map Variable (Indices a)
-> Map
     Variable
     (ReaderT
        (SIEnv a)
        (WriterT EvalLog StencilsAnalysis)
        (Maybe (Bool, [[Int]])))
mkOffsets = (Variable
 -> Indices a
 -> ReaderT
      (SIEnv a)
      (WriterT EvalLog StencilsAnalysis)
      (Maybe (Bool, [[Int]])))
-> Map Variable (Indices a)
-> Map
     Variable
     (ReaderT
        (SIEnv a)
        (WriterT EvalLog StencilsAnalysis)
        (Maybe (Bool, [[Int]])))
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (Variable
-> [Neighbour]
-> Indices a
-> ReaderT
     (SIEnv a)
     (WriterT EvalLog StencilsAnalysis)
     (Maybe (Bool, [[Int]]))
forall a.
Data a =>
Variable
-> [Neighbour]
-> Indices a
-> StencilInferer a (Maybe (Bool, [[Int]]))
`indicesToRelativisedOffsets` [Neighbour]
lhs)

{-| genSubscripts
   Takes * a list of blocks representing an RHS
   Returns a map from array variables to indices, and a list of
   nodes that were visited when computing this information -}
genSubscripts
  :: (Data a, Show a, Eq a)
  => [F.Block (FA.Analysis a)]
  -> StencilInferer a (M.Map Variable (Indices a), [Int])
genSubscripts :: [Block (Analysis a)]
-> StencilInferer a (Map Variable (Indices a), [Int])
genSubscripts [Block (Analysis a)]
blocks = do
  FlowsGraph a
flowsGraph <- StencilInferer a (FlowsGraph a)
forall ann. StencilInferer ann (FlowsGraph ann)
getFlowsGraph
  let ([Map Variable (Indices a)]
maps, [Int]
visitedNodes) = State [Int] [Map Variable (Indices a)]
-> [Int] -> ([Map Variable (Indices a)], [Int])
forall s a. State s a -> s -> (a, s)
runState ((Block (Analysis a)
 -> StateT [Int] Identity (Map Variable (Indices a)))
-> [Block (Analysis a)] -> State [Int] [Map Variable (Indices a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool
-> FlowsGraph a
-> Block (Analysis a)
-> StateT [Int] Identity (Map Variable (Indices a))
forall a.
(Data a, Show a, Eq a) =>
Bool
-> FlowsGraph a
-> Block (Analysis a)
-> State [Int] (Map Variable (Indices a))
genSubscripts' Bool
True FlowsGraph a
flowsGraph) [Block (Analysis a)]
blocks) []
      subscripts :: Map Variable (Indices a)
subscripts = (Indices a -> Indices a -> Indices a)
-> [Map Variable (Indices a)] -> Map Variable (Indices a)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith Indices a -> Indices a -> Indices a
forall a. [a] -> [a] -> [a]
(++) [Map Variable (Indices a)]
maps
  (Map Variable (Indices a), [Int])
-> StencilInferer a (Map Variable (Indices a), [Int])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map Variable (Indices a)
subscripts, [Int]
visitedNodes)
  where
    -- Generate all subscripting expressions (that are translations on
    -- induction variables) that flow to this block
    -- The State monad provides a list of the visited nodes so far
    genSubscripts'
      :: (Data a, Show a, Eq a)
      => Bool
      -> FAD.FlowsGraph a
      -> F.Block (FA.Analysis a)
      -> State [Int] (M.Map Variable (Indices a))

    genSubscripts' :: Bool
-> FlowsGraph a
-> Block (Analysis a)
-> State [Int] (Map Variable (Indices a))
genSubscripts' Bool
False FlowsGraph a
_ (F.BlStatement Analysis a
_ SrcSpan
_ Maybe (Expression (Analysis a))
_ (F.StExpressionAssign Analysis a
_ SrcSpan
_ Expression (Analysis a)
e Expression (Analysis a)
_))
       | Maybe [Index (Analysis a)] -> Bool
forall a. Maybe a -> Bool
isJust (Maybe [Index (Analysis a)] -> Bool)
-> Maybe [Index (Analysis a)] -> Bool
forall a b. (a -> b) -> a -> b
$ Expression (Analysis a) -> Maybe [Index (Analysis a)]
forall a. Expression (Analysis a) -> Maybe [Index (Analysis a)]
isArraySubscript Expression (Analysis a)
e
       -- Don't pull dependencies through arrays
       = Map Variable (Indices a) -> State [Int] (Map Variable (Indices a))
forall (m :: * -> *) a. Monad m => a -> m a
return Map Variable (Indices a)
forall k a. Map k a
M.empty

    genSubscripts' Bool
_ FlowsGraph a
flowsGraph Block (Analysis a)
block = do
       [Int]
visited <- StateT [Int] Identity [Int]
forall s (m :: * -> *). MonadState s m => m s
get
       case Analysis a -> Maybe Int
forall a. Analysis a -> Maybe Int
FA.insLabel (Analysis a -> Maybe Int) -> Analysis a -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Block (Analysis a) -> Analysis a
forall (f :: * -> *) a. Annotated f => f a -> a
F.getAnnotation Block (Analysis a)
block of

         Just Int
node
           | Int
node Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
visited ->
            -- This dependency has already been visited during this traversal
              Map Variable (Indices a) -> State [Int] (Map Variable (Indices a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map Variable (Indices a)
forall k a. Map k a
M.empty
           | Bool
otherwise -> do
            -- Fresh dependency
            [Int] -> StateT [Int] Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([Int] -> StateT [Int] Identity ())
-> [Int] -> StateT [Int] Identity ()
forall a b. (a -> b) -> a -> b
$ Int
node Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
visited
            let blocksFlowingIn :: [Block (Analysis a)]
blocksFlowingIn = (Int -> Maybe (Block (Analysis a)))
-> [Int] -> [Block (Analysis a)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (FlowsGraph a -> Int -> Maybe (Block (Analysis a))
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Maybe a
lab FlowsGraph a
flowsGraph) ([Int] -> [Block (Analysis a)]) -> [Int] -> [Block (Analysis a)]
forall a b. (a -> b) -> a -> b
$ FlowsGraph a -> Int -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> Int -> [Int]
pre FlowsGraph a
flowsGraph Int
node
            -- Try to get the block from the flowsGraph before analysis its rhses
            let blockG :: Block (Analysis a)
blockG = Block (Analysis a)
-> Maybe (Block (Analysis a)) -> Block (Analysis a)
forall a. a -> Maybe a -> a
fromMaybe Block (Analysis a)
block (FlowsGraph a -> Int -> Maybe (Block (Analysis a))
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> Maybe a
lab FlowsGraph a
flowsGraph Int
node)
            [Map Variable (Indices a)]
dependencies <- (Block (Analysis a) -> State [Int] (Map Variable (Indices a)))
-> [Block (Analysis a)]
-> StateT [Int] Identity [Map Variable (Indices a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool
-> FlowsGraph a
-> Block (Analysis a)
-> State [Int] (Map Variable (Indices a))
forall a.
(Data a, Show a, Eq a) =>
Bool
-> FlowsGraph a
-> Block (Analysis a)
-> State [Int] (Map Variable (Indices a))
genSubscripts' Bool
False FlowsGraph a
flowsGraph) [Block (Analysis a)]
blocksFlowingIn
            Map Variable (Indices a) -> State [Int] (Map Variable (Indices a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Map Variable (Indices a)
 -> State [Int] (Map Variable (Indices a)))
-> Map Variable (Indices a)
-> State [Int] (Map Variable (Indices a))
forall a b. (a -> b) -> a -> b
$ (Indices a -> Indices a -> Indices a)
-> [Map Variable (Indices a)] -> Map Variable (Indices a)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith Indices a -> Indices a -> Indices a
forall a. [a] -> [a] -> [a]
(++) (Block (Analysis a) -> Map Variable (Indices a)
forall a.
(Data a, Eq a) =>
Block (Analysis a) -> Map Variable (Indices a)
genRHSsubscripts Block (Analysis a)
blockG Map Variable (Indices a)
-> [Map Variable (Indices a)] -> [Map Variable (Indices a)]
forall a. a -> [a] -> [a]
: [Map Variable (Indices a)]
dependencies)

         Maybe Int
Nothing -> Variable -> State [Int] (Map Variable (Indices a))
forall a. HasCallStack => Variable -> a
error (Variable -> State [Int] (Map Variable (Indices a)))
-> Variable -> State [Int] (Map Variable (Indices a))
forall a b. (a -> b) -> a -> b
$ Variable
"Missing a label for: " Variable -> ShowS
forall a. [a] -> [a] -> [a]
++ Block (Analysis a) -> Variable
forall a. Show a => a -> Variable
show Block (Analysis a)
block

-- | Given an induction variable map, and a piece of syntax
-- return a list of induction variables in scope for this index
extractRelevantIVS :: (FU.Spanned (ast (FA.Analysis a)), F.Annotated ast) =>
     FAD.InductionVarMapByASTBlock
  -> ast (FA.Analysis a)
  -> [Variable]
extractRelevantIVS :: InductionVarMapByASTBlock -> ast (Analysis a) -> [Variable]
extractRelevantIVS InductionVarMapByASTBlock
ivmap ast (Analysis a)
f = [Variable]
ivsList
  where
    ivsList :: [Variable]
ivsList = Set Variable -> [Variable]
forall a. Set a -> [a]
S.toList (Set Variable -> [Variable]) -> Set Variable -> [Variable]
forall a b. (a -> b) -> a -> b
$ Set Variable -> Maybe (Set Variable) -> Set Variable
forall a. a -> Maybe a -> a
fromMaybe Set Variable
forall a. Set a
S.empty (Maybe (Set Variable) -> Set Variable)
-> Maybe (Set Variable) -> Set Variable
forall a b. (a -> b) -> a -> b
$ Int -> InductionVarMapByASTBlock -> Maybe (Set Variable)
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
label InductionVarMapByASTBlock
ivmap

    label :: Int
label   = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Variable -> Int
forall a. HasCallStack => Variable -> a
error Variable
errorMsg) (Analysis a -> Maybe Int
forall a. Analysis a -> Maybe Int
FA.insLabel (Analysis a -> Maybe Int)
-> (ast (Analysis a) -> Analysis a)
-> ast (Analysis a)
-> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ast (Analysis a) -> Analysis a
forall (f :: * -> *) a. Annotated f => f a -> a
F.getAnnotation (ast (Analysis a) -> Maybe Int) -> ast (Analysis a) -> Maybe Int
forall a b. (a -> b) -> a -> b
$ ast (Analysis a)
f)
    -- For debugging purposes
    errorMsg :: Variable
errorMsg = SrcSpan -> Variable
forall a. Show a => a -> Variable
show (ast (Analysis a) -> SrcSpan
forall a. Spanned a => a -> SrcSpan
FU.getSpan ast (Analysis a)
f)
            Variable -> ShowS
forall a. [a] -> [a] -> [a]
++ Variable
" get IVs associated to labelled index "

{-| Given a list of induction variables and an index, compute
   its Neighbour representation
   e.g., for the expression a(i+1,j-1) then this function gets
   passed expr = i + 1   (returning +1) and expr = j - 1 (returning -1) -}
convIxToNeighbour :: (Data a) => [Variable] -> F.Index (FA.Analysis a) -> Neighbour
convIxToNeighbour :: [Variable] -> Index (Analysis a) -> Neighbour
convIxToNeighbour [Variable]
_ (F.IxRange Analysis a
_ SrcSpan
_ Maybe (Expression (Analysis a))
Nothing Maybe (Expression (Analysis a))
Nothing Maybe (Expression (Analysis a))
Nothing)     = Variable -> Int -> Neighbour
Neighbour Variable
"" Int
0
convIxToNeighbour [Variable]
_ (F.IxRange Analysis a
_ SrcSpan
_ Maybe (Expression (Analysis a))
Nothing Maybe (Expression (Analysis a))
Nothing
                  (Just (F.ExpValue Analysis a
_ SrcSpan
_ (F.ValInteger Variable
"1")))) = Variable -> Int -> Neighbour
Neighbour Variable
"" Int
0

convIxToNeighbour [Variable]
ivs (F.IxSingle Analysis a
_ SrcSpan
_ Maybe Variable
_ Expression (Analysis a)
expr)  = [Variable] -> Expression (Analysis a) -> Neighbour
forall a.
Data a =>
[Variable] -> Expression (Analysis a) -> Neighbour
expToNeighbour [Variable]
ivs Expression (Analysis a)
expr
convIxToNeighbour [Variable]
_ Index (Analysis a)
_ = Neighbour
NonNeighbour -- indexing expression is a range

-- Combinator for reducing a map with effects and partiality inside
-- into an effectful list of key-value pairs
assocsSequence :: Monad m => M.Map k (m (Maybe a)) -> m [(k, a)]
assocsSequence :: Map k (m (Maybe a)) -> m [(k, a)]
assocsSequence Map k (m (Maybe a))
maps = do
  [(k, Maybe a)]
assocs <- ((k, m (Maybe a)) -> m (k, Maybe a))
-> [(k, m (Maybe a))] -> m [(k, Maybe a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (k, m (Maybe a)) -> m (k, Maybe a)
forall (m :: * -> *) a b. Monad m => (a, m b) -> m (a, b)
strength ([(k, m (Maybe a))] -> m [(k, Maybe a)])
-> (Map k (m (Maybe a)) -> [(k, m (Maybe a))])
-> Map k (m (Maybe a))
-> m [(k, Maybe a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k (m (Maybe a)) -> [(k, m (Maybe a))]
forall k a. Map k a -> [(k, a)]
M.toList (Map k (m (Maybe a)) -> m [(k, Maybe a)])
-> Map k (m (Maybe a)) -> m [(k, Maybe a)]
forall a b. (a -> b) -> a -> b
$ Map k (m (Maybe a))
maps
  [(k, a)] -> m [(k, a)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(k, a)] -> m [(k, a)])
-> ([(k, Maybe a)] -> [(k, a)]) -> [(k, Maybe a)] -> m [(k, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, Maybe a) -> Maybe (k, a)) -> [(k, Maybe a)] -> [(k, a)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (k, Maybe a) -> Maybe (k, a)
forall (m :: * -> *) a b. Monad m => (a, m b) -> m (a, b)
strength ([(k, Maybe a)] -> m [(k, a)]) -> [(k, Maybe a)] -> m [(k, a)]
forall a b. (a -> b) -> a -> b
$ [(k, Maybe a)]
assocs
  where
    strength :: Monad m => (a, m b) -> m (a, b)
    strength :: (a, m b) -> m (a, b)
strength (a
a, m b
mb) = m b
mb m b -> (b -> m (a, b)) -> m (a, b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\b
b -> (a, b) -> m (a, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, b
b))

-- Convert list of indexing expressions to a spec
indicesToSpec :: (Data a)
              => Variable
              -> [Neighbour]
              -> Indices a
              -> StencilInferer a (Maybe Specification)
indicesToSpec :: Variable
-> [Neighbour]
-> Indices a
-> StencilInferer a (Maybe Specification)
indicesToSpec Variable
a [Neighbour]
lhs Indices a
ixs = do
  Maybe (Bool, [[Int]])
mMultOffsets <- Variable
-> [Neighbour]
-> Indices a
-> StencilInferer a (Maybe (Bool, [[Int]]))
forall a.
Data a =>
Variable
-> [Neighbour]
-> Indices a
-> StencilInferer a (Maybe (Bool, [[Int]]))
indicesToRelativisedOffsets Variable
a [Neighbour]
lhs Indices a
ixs
  Maybe Specification -> StencilInferer a (Maybe Specification)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Specification -> StencilInferer a (Maybe Specification))
-> Maybe Specification -> StencilInferer a (Maybe Specification)
forall a b. (a -> b) -> a -> b
$ do
    (Bool
mult, [[Int]]
offsets) <- Maybe (Bool, [[Int]])
mMultOffsets
    Specification
spec <- [[Int]] -> Maybe Specification
relativeIxsToSpec [[Int]]
offsets
    let spec' :: Specification
spec' = Linearity -> Specification -> Specification
setLinearity (Bool -> Linearity
fromBool Bool
mult) Specification
spec
    Specification -> Maybe Specification
forall (m :: * -> *) a. Monad m => a -> m a
return (Specification -> Maybe Specification)
-> Specification -> Maybe Specification
forall a b. (a -> b) -> a -> b
$ [Neighbour] -> Specification -> Specification
setType [Neighbour]
lhs Specification
spec'

-- Get all RHS subscript which are translated induction variables
-- return as a map from (source name) variables to a list of relative indices
genRHSsubscripts :: forall a. (Data a, Eq a)
                 => F.Block (FA.Analysis a) -> M.Map Variable (Indices a)
genRHSsubscripts :: Block (Analysis a) -> Map Variable (Indices a)
genRHSsubscripts Block (Analysis a)
block = Block (Analysis a) -> Map Variable (Indices a)
forall a (b :: * -> *).
(Eq a, Data a, Data (b (Analysis a))) =>
b (Analysis a) -> Map Variable [[Index (Analysis a)]]
genRHSsubscripts' ((Expression (Analysis a) -> Expression (Analysis a))
-> Block (Analysis a) -> Block (Analysis a)
forall from to. Biplate from to => (to -> to) -> from -> from
transformBi Expression (Analysis a) -> Expression (Analysis a)
replaceModulo Block (Analysis a)
block)
  where
    -- Any occurence of an subscript "modulo(e, e')" is replaced with "e"
    replaceModulo :: F.Expression (FA.Analysis a) -> F.Expression (FA.Analysis a)
    replaceModulo :: Expression (Analysis a) -> Expression (Analysis a)
replaceModulo (F.ExpFunctionCall Analysis a
_ SrcSpan
_
                      (F.ExpValue Analysis a
_ SrcSpan
_ (F.ValIntrinsic Variable
iname)) Maybe (AList Argument (Analysis a))
subs)
        | Variable
iname Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable
"modulo", Variable
"mod", Variable
"amod", Variable
"dmod"]
        -- We expect that the first parameter to modulo is being treated
        -- as an IxSingle element
        , Just (F.Argument Analysis a
_ SrcSpan
_ Maybe Variable
_ Expression (Analysis a)
e':[Argument (Analysis a)]
_) <- (AList Argument (Analysis a) -> [Argument (Analysis a)])
-> Maybe (AList Argument (Analysis a))
-> Maybe [Argument (Analysis a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AList Argument (Analysis a) -> [Argument (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
F.aStrip Maybe (AList Argument (Analysis a))
subs = Expression (Analysis a)
e'
    replaceModulo Expression (Analysis a)
e = Expression (Analysis a)
e

    genRHSsubscripts' :: b (Analysis a) -> Map Variable [[Index (Analysis a)]]
genRHSsubscripts' b (Analysis a)
b =
       [(Variable, [Index (Analysis a)])]
-> Map Variable [[Index (Analysis a)]]
forall a k. (Eq a, Ord k) => [(k, a)] -> Map k [a]
collect [ (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.srcName Expression (Analysis a)
expr, [Index (Analysis a)]
e)
         | F.ExpSubscript Analysis a
_ SrcSpan
_ Expression (Analysis a)
expr AList Index (Analysis a)
subs <- b (Analysis a) -> [Expression (Analysis a)]
forall a (b :: * -> *).
(Data a, Data (b a)) =>
b a -> [Expression a]
FA.rhsExprs b (Analysis a)
b
         , Expression (Analysis a) -> Bool
forall a. Expression a -> Bool
isVariableExpr Expression (Analysis a)
expr
         , let e :: [Index (Analysis a)]
e = AList Index (Analysis a) -> [Index (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
F.aStrip AList Index (Analysis a)
subs
         , Bool -> Bool
not ([Index (Analysis a)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Index (Analysis a)]
e)]

-- Given a list of induction variables and an expression, compute its
-- Neighbour representation
expToNeighbour :: forall a. Data a
            => [Variable] -> F.Expression (FA.Analysis a) -> Neighbour

expToNeighbour :: [Variable] -> Expression (Analysis a) -> Neighbour
expToNeighbour [Variable]
ivs e :: Expression (Analysis a)
e@(F.ExpValue Analysis a
_ SrcSpan
_ v :: Value (Analysis a)
v@(F.ValVariable Variable
_))
    | Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs = Variable -> Int -> Neighbour
Neighbour (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e) Int
0
    | Bool
otherwise               = Value () -> Neighbour
Constant (Value (Analysis a) -> Value ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void Value (Analysis a)
v)

expToNeighbour [Variable]
_ (F.ExpValue Analysis a
_ SrcSpan
_ Value (Analysis a)
val) = Value () -> Neighbour
Constant (Value (Analysis a) -> Value ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void Value (Analysis a)
val)

-- use constant-expression analysis if available
expToNeighbour [Variable]
ivs (F.ExpBinary Analysis a
_ SrcSpan
_ BinaryOp
F.Addition
                    e1 :: Expression (Analysis a)
e1@(F.ExpValue Analysis a
_ SrcSpan
_ (F.ValVariable Variable
_))
                    Expression (Analysis a)
e2)
    | Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e1 Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs
    , Just (FAD.ConstInt Integer
offs) <- Analysis a -> Maybe Constant
forall a. Analysis a -> Maybe Constant
FA.constExp (Expression (Analysis a) -> Analysis a
forall (f :: * -> *) a. Annotated f => f a -> a
F.getAnnotation Expression (Analysis a)
e2) = Variable -> Int -> Neighbour
Neighbour (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e1) (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
offs)

expToNeighbour [Variable]
ivs (F.ExpBinary Analysis a
_ SrcSpan
_ BinaryOp
F.Addition
                 e :: Expression (Analysis a)
e@(F.ExpValue Analysis a
_ SrcSpan
_ (F.ValVariable Variable
_))
                   (F.ExpValue Analysis a
_ SrcSpan
_ (F.ValInteger Variable
offs)))
    | Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs = Variable -> Int -> Neighbour
Neighbour (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e) (Variable -> Int
forall a. Read a => Variable -> a
read Variable
offs)

-- use constant-expression analysis if available
expToNeighbour [Variable]
ivs (F.ExpBinary Analysis a
_ SrcSpan
_ BinaryOp
F.Addition
                    e1 :: Expression (Analysis a)
e1@(F.ExpValue Analysis a
_ SrcSpan
_ (F.ValVariable Variable
_))
                    Expression (Analysis a)
e2)
    | Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e1 Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs
    , Just (FAD.ConstInt Integer
offs) <- Analysis a -> Maybe Constant
forall a. Analysis a -> Maybe Constant
FA.constExp (Expression (Analysis a) -> Analysis a
forall (f :: * -> *) a. Annotated f => f a -> a
F.getAnnotation Expression (Analysis a)
e2) = Variable -> Int -> Neighbour
Neighbour (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e1) (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
offs)

expToNeighbour [Variable]
ivs (F.ExpBinary Analysis a
_ SrcSpan
_ BinaryOp
F.Addition
                  (F.ExpValue Analysis a
_ SrcSpan
_ (F.ValInteger Variable
offs))
                e :: Expression (Analysis a)
e@(F.ExpValue Analysis a
_ SrcSpan
_ (F.ValVariable Variable
_)))
    | Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs = Variable -> Int -> Neighbour
Neighbour (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e) (Variable -> Int
forall a. Read a => Variable -> a
read Variable
offs)

-- use constant-expression analysis if available
expToNeighbour [Variable]
ivs (F.ExpBinary Analysis a
_ SrcSpan
_ BinaryOp
F.Subtraction
                    e1 :: Expression (Analysis a)
e1@(F.ExpValue Analysis a
_ SrcSpan
_ (F.ValVariable Variable
_))
                    Expression (Analysis a)
e2)
   | Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e1 Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs
   , Just (FAD.ConstInt Integer
offs) <- Analysis a -> Maybe Constant
forall a. Analysis a -> Maybe Constant
FA.constExp (Expression (Analysis a) -> Analysis a
forall (f :: * -> *) a. Annotated f => f a -> a
F.getAnnotation Expression (Analysis a)
e2)
   , Integer
offs' <- if Integer
offs Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 then Integer -> Integer
forall a. Num a => a -> a
abs Integer
offs else (- Integer
offs) = Variable -> Int -> Neighbour
Neighbour (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e1) (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
offs')

expToNeighbour [Variable]
ivs (F.ExpBinary Analysis a
_ SrcSpan
_ BinaryOp
F.Subtraction
                 e :: Expression (Analysis a)
e@(F.ExpValue Analysis a
_ SrcSpan
_ (F.ValVariable Variable
_))
                   (F.ExpValue Analysis a
_ SrcSpan
_ (F.ValInteger Variable
offs)))
   | Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs =
         Variable -> Int -> Neighbour
Neighbour (Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e) (if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Int -> Int
forall a. Num a => a -> a
abs Int
x else (- Int
x))
             where x :: Int
x = Variable -> Int
forall a. Read a => Variable -> a
read Variable
offs

expToNeighbour [Variable]
ivs Expression (Analysis a)
expr =
  -- Record when there is some kind of relative index on an inducion variable
  -- but that is not a neighbourhood index by our definitions
  if [Variable] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Variable]
ivs' then Value () -> Neighbour
Constant (Variable -> Value ()
forall a. Variable -> Value a
F.ValInteger Variable
"0") else Neighbour
NonNeighbour
  where
    -- set of all induction variables involved in this expression
    ivs' :: [Variable]
ivs' = [Variable
i | e :: Expression (Analysis a)
e@(F.ExpValue Analysis a
_ SrcSpan
_ F.ValVariable{})
                 <- Expression (Analysis a) -> [Expression (Analysis a)]
forall from to. Biplate from to => from -> [to]
universeBi Expression (Analysis a)
expr :: [F.Expression (FA.Analysis a)]
                , let i :: Variable
i = Expression (Analysis a) -> Variable
forall a. Expression (Analysis a) -> Variable
FA.varName Expression (Analysis a)
e
                , Variable
i Variable -> [Variable] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Variable]
ivs]

indicesToRelativisedOffsets :: (Data a)
                            => Variable
                            -> [Neighbour]
                            -> Indices a
                            -> StencilInferer a (Maybe (Bool, [[Int]]))
indicesToRelativisedOffsets :: Variable
-> [Neighbour]
-> Indices a
-> StencilInferer a (Maybe (Bool, [[Int]]))
indicesToRelativisedOffsets Variable
a [Neighbour]
lhs Indices a
ixs = do
  [Variable]
ivs <- StencilInferer a [Variable]
forall ann. StencilInferer ann [Variable]
getIvs
   -- Convert indices to neighbourhood representation
  let rhses :: [[Neighbour]]
rhses = ([Index (Analysis a)] -> [Neighbour]) -> Indices a -> [[Neighbour]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Index (Analysis a) -> Neighbour)
-> [Index (Analysis a)] -> [Neighbour]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Variable] -> Index (Analysis a) -> Neighbour
forall a. Data a => [Variable] -> Index (Analysis a) -> Neighbour
convIxToNeighbour [Variable]
ivs)) Indices a
ixs

  -- As an optimisation, do duplicate check in front-end first
  -- so that duplicate indices don't get passed into the main engine
  let ([[Neighbour]]
rhses', Bool
mult) = [[Neighbour]] -> ([[Neighbour]], Bool)
forall a. Eq a => [a] -> ([a], Bool)
hasDuplicates [[Neighbour]]
rhses

  -- Check that induction variables are used consistently on lhs and rhses
  if Bool -> Bool
not ([Neighbour] -> [[Neighbour]] -> Bool
consistentIVSuse [Neighbour]
lhs [[Neighbour]]
rhses')
    then do EvalLog -> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [(Variable
"EVALMODE: Inconsistent IV use (tag: inconsistentIV)", Variable
"")]
            Maybe (Bool, [[Int]]) -> StencilInferer a (Maybe (Bool, [[Int]]))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Bool, [[Int]])
forall a. Maybe a
Nothing
    else
      -- For the EvalMode, if there are any non-neighbourhood relative
      -- subscripts detected then add this to the eval log
      if [[Neighbour]] -> Bool
hasNonNeighbourhoodRelatives [[Neighbour]]
rhses'
      then do EvalLog -> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [(Variable
"EVALMODE: Non-neighbour relative subscripts\
                    \ (tag: nonNeighbour)",Variable
"")]
              Maybe (Bool, [[Int]]) -> StencilInferer a (Maybe (Bool, [[Int]]))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Bool, [[Int]])
forall a. Maybe a
Nothing
      else do
        -- Relativize the offsets based on the lhs
        let rhses'' :: [[Neighbour]]
rhses'' = [Neighbour] -> [[Neighbour]] -> [[Neighbour]]
relativise [Neighbour]
lhs [[Neighbour]]
rhses'
        Bool
-> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
-> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([[Neighbour]]
rhses' [[Neighbour]] -> [[Neighbour]] -> Bool
forall a. Eq a => a -> a -> Bool
/= [[Neighbour]]
rhses'') (ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
 -> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ())
-> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
-> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
forall a b. (a -> b) -> a -> b
$
          EvalLog -> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [(Variable
"EVALMODE: Relativized spec (tag: relativized)", Variable
"")]

        let offsets :: [[Int]]
offsets  = [[Int]] -> [[Int]]
padZeros ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ ([Neighbour] -> [Int]) -> [[Neighbour]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe [Int] -> [Int]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [Int] -> [Int])
-> ([Neighbour] -> Maybe [Int]) -> [Neighbour] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Neighbour -> Maybe Int) -> [Neighbour] -> Maybe [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Neighbour -> Maybe Int
neighbourToOffset) [[Neighbour]]
rhses''
        EvalLog -> ReaderT (SIEnv a) (WriterT EvalLog StencilsAnalysis) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [(Variable
"EVALMODE: dimensionality=" Variable -> ShowS
forall a. [a] -> [a] -> [a]
++
                 Int -> Variable
forall a. Show a => a -> Variable
show (if [[Int]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Int]]
offsets then Int
0 else [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> Int) -> ([[Int]] -> [Int]) -> [[Int]] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Int]] -> [Int]
forall a. [a] -> a
head ([[Int]] -> Int) -> [[Int]] -> Int
forall a b. (a -> b) -> a -> b
$ [[Int]]
offsets), Variable
a)]
        Maybe (Bool, [[Int]]) -> StencilInferer a (Maybe (Bool, [[Int]]))
forall (m :: * -> *) a. Monad m => a -> m a
return ((Bool, [[Int]]) -> Maybe (Bool, [[Int]])
forall a. a -> Maybe a
Just (Bool
mult, [[Int]]
offsets))
  where hasNonNeighbourhoodRelatives :: [[Neighbour]] -> Bool
hasNonNeighbourhoodRelatives = ([Neighbour] -> Bool) -> [[Neighbour]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Neighbour -> [Neighbour] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Neighbour
NonNeighbour)

-- Convert list of relative offsets to a spec
relativeIxsToSpec :: [[Int]] -> Maybe Specification
relativeIxsToSpec :: [[Int]] -> Maybe Specification
relativeIxsToSpec [[Int]]
ixs =
  if Specification -> Bool
isEmpty Specification
exactSpec then Maybe Specification
forall a. Maybe a
Nothing else Specification -> Maybe Specification
forall a. a -> Maybe a
Just Specification
exactSpec
  where exactSpec :: Specification
exactSpec = VecList Int -> Specification
inferFromIndicesWithoutLinearity (VecList Int -> Specification)
-> ([[Int]] -> VecList Int) -> [[Int]] -> Specification
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Int]] -> VecList Int
forall a. [[a]] -> VecList a
V.fromLists ([[Int]] -> Specification) -> [[Int]] -> Specification
forall a b. (a -> b) -> a -> b
$ [[Int]]
ixs

{-| Set the type of Specification (stencil or access) based on the lhs
    set of neighbourhood indices; empty implies this is an access
    specification -}
setType :: [Neighbour] -> Specification -> Specification
setType :: [Neighbour] -> Specification -> Specification
setType [] (Specification Multiplicity (Approximation Spatial)
spec Bool
_) = Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification Multiplicity (Approximation Spatial)
spec Bool
False
setType [Neighbour]
xs (Specification Multiplicity (Approximation Spatial)
spec Bool
_) | (Neighbour -> Bool) -> [Neighbour] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Neighbour -> Bool
isConstant [Neighbour]
xs = Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification Multiplicity (Approximation Spatial)
spec Bool
False
  where
    isConstant :: Neighbour -> Bool
isConstant (Constant Value ()
_) = Bool
True
    isConstant Neighbour
_            = Bool
False
setType [Neighbour]
_  (Specification Multiplicity (Approximation Spatial)
spec Bool
_)  = Multiplicity (Approximation Spatial) -> Bool -> Specification
Specification Multiplicity (Approximation Spatial)
spec Bool
True

-- Given a list of the neighbourhood representation for the LHS, of size n
-- and a list of size-n lists of offsets, relativise the offsets
relativise :: [Neighbour] -> [[Neighbour]] -> [[Neighbour]]
relativise :: [Neighbour] -> [[Neighbour]] -> [[Neighbour]]
relativise [Neighbour]
lhs [[Neighbour]]
rhses = (Neighbour -> [[Neighbour]] -> [[Neighbour]])
-> [[Neighbour]] -> [Neighbour] -> [[Neighbour]]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Neighbour -> [[Neighbour]] -> [[Neighbour]]
relativiseRHS [[Neighbour]]
rhses [Neighbour]
lhs
  where
    relativiseRHS :: Neighbour -> [[Neighbour]] -> [[Neighbour]]
relativiseRHS (Neighbour Variable
lhsIV Int
i) [[Neighbour]]
rs =
        ([Neighbour] -> [Neighbour]) -> [[Neighbour]] -> [[Neighbour]]
forall a b. (a -> b) -> [a] -> [b]
map ((Neighbour -> Neighbour) -> [Neighbour] -> [Neighbour]
forall a b. (a -> b) -> [a] -> [b]
map (Variable -> Int -> Neighbour -> Neighbour
relativiseBy Variable
lhsIV Int
i)) [[Neighbour]]
rs
    relativiseRHS Neighbour
_ [[Neighbour]]
rs = [[Neighbour]]
rs

    relativiseBy :: Variable -> Int -> Neighbour -> Neighbour
relativiseBy Variable
v Int
i (Neighbour Variable
u Int
j) | Variable
v Variable -> Variable -> Bool
forall a. Eq a => a -> a -> Bool
== Variable
u = Variable -> Int -> Neighbour
Neighbour Variable
u (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)
    relativiseBy Variable
_ Int
_ Neighbour
x = Neighbour
x

-- Helper predicates
isVariableExpr :: F.Expression a -> Bool
isVariableExpr :: Expression a -> Bool
isVariableExpr (F.ExpValue a
_ SrcSpan
_ (F.ValVariable Variable
_)) = Bool
True
isVariableExpr Expression a
_                                  = Bool
False

-- Check that induction variables are used consistently
consistentIVSuse :: [Neighbour] -> [[Neighbour]] -> Bool
consistentIVSuse :: [Neighbour] -> [[Neighbour]] -> Bool
consistentIVSuse [] [[Neighbour]]
_ = Bool
True
consistentIVSuse [Neighbour]
_ [] = Bool
True
consistentIVSuse [Neighbour]
lhs [[Neighbour]]
rhses =
     Maybe [Neighbour] -> Bool
forall a. Maybe a -> Bool
isJust Maybe [Neighbour]
rhsBasis -- There is a consitent RHS
  Bool -> Bool -> Bool
&& ((Neighbour -> Bool) -> [Neighbour] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Neighbour -> [Neighbour] -> Bool
`consistentWith` [Neighbour]
lhs) (Maybe [Neighbour] -> [Neighbour]
forall a. HasCallStack => Maybe a -> a
fromJust Maybe [Neighbour]
rhsBasis)
   Bool -> Bool -> Bool
|| (Neighbour -> Bool) -> [Neighbour] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Neighbour -> [Neighbour] -> Bool
`consistentWith` Maybe [Neighbour] -> [Neighbour]
forall a. HasCallStack => Maybe a -> a
fromJust Maybe [Neighbour]
rhsBasis) [Neighbour]
lhs)
    where
      cmp :: Neighbour -> Neighbour -> Maybe Neighbour
cmp (Neighbour Variable
v Int
i) (Neighbour Variable
v' Int
_) | Variable
v Variable -> Variable -> Bool
forall a. Eq a => a -> a -> Bool
== Variable
v'   = Neighbour -> Maybe Neighbour
forall a. a -> Maybe a
Just (Neighbour -> Maybe Neighbour) -> Neighbour -> Maybe Neighbour
forall a b. (a -> b) -> a -> b
$ Variable -> Int -> Neighbour
Neighbour Variable
v Int
i
                                           | Bool
otherwise = Maybe Neighbour
forall a. Maybe a
Nothing
      -- Cases for constants or non neighbour indices
      cmp n :: Neighbour
n@Neighbour{}  (Constant Value ()
_) = Neighbour -> Maybe Neighbour
forall a. a -> Maybe a
Just Neighbour
n
      cmp (Constant Value ()
_) n :: Neighbour
n@Neighbour{}  = Neighbour -> Maybe Neighbour
forall a. a -> Maybe a
Just Neighbour
n
      cmp NonNeighbour{} Neighbour{}  = Maybe Neighbour
forall a. Maybe a
Nothing
      cmp Neighbour{} NonNeighbour{}  = Maybe Neighbour
forall a. Maybe a
Nothing
      cmp Neighbour
_ Neighbour
_                         = Neighbour -> Maybe Neighbour
forall a. a -> Maybe a
Just (Neighbour -> Maybe Neighbour) -> Neighbour -> Maybe Neighbour
forall a b. (a -> b) -> a -> b
$ Value () -> Neighbour
Constant (Variable -> Value ()
forall a. Variable -> Value a
F.ValInteger Variable
"")
      rhsBasis :: Maybe [Neighbour]
rhsBasis = ([Neighbour] -> [Neighbour] -> Maybe [Neighbour])
-> [Neighbour] -> [[Neighbour]] -> Maybe [Neighbour]
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM ((Neighbour -> Neighbour -> Maybe Neighbour)
-> [Neighbour] -> [Neighbour] -> Maybe [Neighbour]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Neighbour -> Neighbour -> Maybe Neighbour
cmp) ([[Neighbour]] -> [Neighbour]
forall a. [a] -> a
head [[Neighbour]]
rhses) ([[Neighbour]] -> [[Neighbour]]
forall a. [a] -> [a]
tail [[Neighbour]]
rhses)
      -- If there is an induction variable on the RHS, then it also occurs on
      -- the LHS
      consistentWith :: Neighbour -> [Neighbour] -> Bool
      consistentWith :: Neighbour -> [Neighbour] -> Bool
consistentWith (Neighbour Variable
rv Int
_) [Neighbour]
ns = (Neighbour -> Bool) -> [Neighbour] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Variable -> Neighbour -> Bool
matchesIV Variable
rv) [Neighbour]
ns
      consistentWith Neighbour
_                [Neighbour]
_  = Bool
True

      matchesIV :: Variable -> Neighbour -> Bool
      matchesIV :: Variable -> Neighbour -> Bool
matchesIV Variable
v (Neighbour Variable
v' Int
_) | Variable
v Variable -> Variable -> Bool
forall a. Eq a => a -> a -> Bool
== Variable
v' = Bool
True
      -- All RHS to contain index ranges
      matchesIV Variable
v Neighbour{}      | Variable
v  Variable -> Variable -> Bool
forall a. Eq a => a -> a -> Bool
== Variable
"" = Bool
True
      matchesIV Variable
_ (Neighbour Variable
v' Int
_) | Variable
v' Variable -> Variable -> Bool
forall a. Eq a => a -> a -> Bool
== Variable
"" = Bool
True
      matchesIV Variable
_ Neighbour
_                          = Bool
False

-- padZeros makes this rectilinear
padZeros :: [[Int]] -> [[Int]]
padZeros :: [[Int]] -> [[Int]]
padZeros [[Int]]
ixss = let m :: Int
m = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
ixss)
                in ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (\[Int]
ixs -> [Int]
ixs [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ixs) Int
0) [[Int]]
ixss

neighbourToOffset :: Neighbour -> Maybe Int
neighbourToOffset :: Neighbour -> Maybe Int
neighbourToOffset (Neighbour Variable
_ Int
o) = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
o
neighbourToOffset (Constant Value ()
_)    = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
absoluteRep
neighbourToOffset Neighbour
_               = Maybe Int
forall a. Maybe a
Nothing