{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Futhark.Representation.SOACS.Simplify
       ( simplifySOACS
       , simplifyLambda
       , simplifyFun
       , simplifyStms
       , simplifyConsts

       , simpleSOACS
       , simplifySOAC

       , soacRules
       )
where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.State
import Control.Monad.Writer
import Data.Foldable
import Data.Either
import Data.List (partition, transpose, unzip6, zip6)
import Data.Maybe
import qualified Data.Map.Strict as M
import qualified Data.Set      as S

import Futhark.Representation.SOACS
import qualified Futhark.Representation.AST as AST
import Futhark.Representation.AST.Attributes.Aliases
import qualified Futhark.Optimise.Simplify.Engine as Engine
import qualified Futhark.Optimise.Simplify as Simplify
import Futhark.Optimise.Simplify.Rules
import Futhark.MonadFreshNames
import Futhark.Optimise.Simplify.Rule
import Futhark.Optimise.Simplify.ClosedForm
import Futhark.Optimise.Simplify.Lore
import Futhark.Tools
import Futhark.Pass
import qualified Futhark.Analysis.SymbolTable as ST
import qualified Futhark.Analysis.UsageTable as UT
import Futhark.Analysis.DataDependencies
import Futhark.Transform.Rename
import Futhark.Util

simpleSOACS :: Simplify.SimpleOps SOACS
simpleSOACS :: SimpleOps SOACS
simpleSOACS = SimplifyOp SOACS (Op SOACS) -> SimpleOps SOACS
forall lore.
(SimplifiableLore lore, Bindable lore) =>
SimplifyOp lore (Op lore) -> SimpleOps lore
Simplify.bindableSimpleOps SimplifyOp SOACS (Op SOACS)
forall lore. SimplifiableLore lore => SimplifyOp lore (SOAC lore)
simplifySOAC

simplifySOACS :: Prog SOACS -> PassM (Prog SOACS)
simplifySOACS :: Prog SOACS -> PassM (Prog SOACS)
simplifySOACS = SimpleOps SOACS
-> RuleBook (Wise SOACS)
-> HoistBlockers SOACS
-> Prog SOACS
-> PassM (Prog SOACS)
forall lore.
SimplifiableLore lore =>
SimpleOps lore
-> RuleBook (Wise lore)
-> HoistBlockers lore
-> Prog lore
-> PassM (Prog lore)
Simplify.simplifyProg SimpleOps SOACS
simpleSOACS RuleBook (Wise SOACS)
soacRules HoistBlockers SOACS
blockers
  where blockers :: HoistBlockers SOACS
blockers = HoistBlockers SOACS
forall lore. HoistBlockers lore
Engine.noExtraHoistBlockers { getArraySizes :: Stm (Wise SOACS) -> Names
Engine.getArraySizes = Stm (Wise SOACS) -> Names
forall lore.
(LetAttr lore ~ (VarWisdom, Type)) =>
Stm lore -> Names
getShapeNames }

-- | Getting the roots of what to hoist, for now only variable
-- names that represent shapes/sizes.
getShapeNames :: (LetAttr lore ~ (VarWisdom, Type)) =>
                 AST.Stm lore -> Names
getShapeNames :: Stm lore -> Names
getShapeNames Stm lore
bnd =
  let tps1 :: [Type]
tps1 = (PatElemT (VarWisdom, Type) -> Type)
-> [PatElemT (VarWisdom, Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT (VarWisdom, Type) -> Type
forall attr. Typed attr => PatElemT attr -> Type
patElemType ([PatElemT (VarWisdom, Type)] -> [Type])
-> [PatElemT (VarWisdom, Type)] -> [Type]
forall a b. (a -> b) -> a -> b
$ PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternElements (PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)])
-> PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall a b. (a -> b) -> a -> b
$ Stm lore -> Pattern lore
forall lore. Stm lore -> Pattern lore
stmPattern Stm lore
bnd
      tps2 :: [Type]
tps2 = (PatElemT (VarWisdom, Type) -> Type)
-> [PatElemT (VarWisdom, Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map ((VarWisdom, Type) -> Type
forall a b. (a, b) -> b
snd ((VarWisdom, Type) -> Type)
-> (PatElemT (VarWisdom, Type) -> (VarWisdom, Type))
-> PatElemT (VarWisdom, Type)
-> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT (VarWisdom, Type) -> (VarWisdom, Type)
forall attr. PatElemT attr -> attr
patElemAttr) ([PatElemT (VarWisdom, Type)] -> [Type])
-> [PatElemT (VarWisdom, Type)] -> [Type]
forall a b. (a -> b) -> a -> b
$ PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternElements (PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)])
-> PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall a b. (a -> b) -> a -> b
$ Stm lore -> Pattern lore
forall lore. Stm lore -> Pattern lore
stmPattern Stm lore
bnd
  in  [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [VName]
subExpVars ([SubExp] -> [VName]) -> [SubExp] -> [VName]
forall a b. (a -> b) -> a -> b
$ (Type -> [SubExp]) -> [Type] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims ([Type]
tps1 [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
tps2)

simplifyFun :: MonadFreshNames m =>
               ST.SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS)
simplifyFun :: SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS)
simplifyFun =
  SimpleOps SOACS
-> RuleBook (Wise SOACS)
-> HoistBlockers SOACS
-> SymbolTable (Wise SOACS)
-> FunDef SOACS
-> m (FunDef SOACS)
forall (m :: * -> *) lore.
(MonadFreshNames m, SimplifiableLore lore) =>
SimpleOps lore
-> RuleBook (Wise lore)
-> HoistBlockers lore
-> SymbolTable (Wise lore)
-> FunDef lore
-> m (FunDef lore)
Simplify.simplifyFun SimpleOps SOACS
simpleSOACS RuleBook (Wise SOACS)
soacRules HoistBlockers SOACS
forall lore. HoistBlockers lore
Engine.noExtraHoistBlockers

simplifyLambda :: (HasScope SOACS m, MonadFreshNames m) =>
                  Lambda -> [Maybe VName] -> m Lambda
simplifyLambda :: Lambda -> [Maybe VName] -> m Lambda
simplifyLambda =
  SimpleOps SOACS
-> RuleBook (Wise SOACS)
-> HoistBlockers SOACS
-> Lambda
-> [Maybe VName]
-> m Lambda
forall (m :: * -> *) lore.
(MonadFreshNames m, HasScope lore m, SimplifiableLore lore) =>
SimpleOps lore
-> RuleBook (Wise lore)
-> HoistBlockers lore
-> Lambda lore
-> [Maybe VName]
-> m (Lambda lore)
Simplify.simplifyLambda SimpleOps SOACS
simpleSOACS RuleBook (Wise SOACS)
soacRules HoistBlockers SOACS
forall lore. HoistBlockers lore
Engine.noExtraHoistBlockers

simplifyStms :: (HasScope SOACS m, MonadFreshNames m) =>
                Stms SOACS -> m (ST.SymbolTable (Wise SOACS), Stms SOACS)
simplifyStms :: Stms SOACS -> m (SymbolTable (Wise SOACS), Stms SOACS)
simplifyStms Stms SOACS
stms = do
  Scope SOACS
scope <- m (Scope SOACS)
forall lore (m :: * -> *). HasScope lore m => m (Scope lore)
askScope
  SimpleOps SOACS
-> RuleBook (Wise SOACS)
-> HoistBlockers SOACS
-> Scope SOACS
-> Stms SOACS
-> m (SymbolTable (Wise SOACS), Stms SOACS)
forall (m :: * -> *) lore.
(MonadFreshNames m, SimplifiableLore lore) =>
SimpleOps lore
-> RuleBook (Wise lore)
-> HoistBlockers lore
-> Scope lore
-> Stms lore
-> m (SymbolTable (Wise lore), Stms lore)
Simplify.simplifyStms SimpleOps SOACS
simpleSOACS RuleBook (Wise SOACS)
soacRules HoistBlockers SOACS
forall lore. HoistBlockers lore
Engine.noExtraHoistBlockers
    Scope SOACS
scope Stms SOACS
stms

simplifyConsts :: MonadFreshNames m =>
                  Stms SOACS -> m (ST.SymbolTable (Wise SOACS), Stms SOACS)
simplifyConsts :: Stms SOACS -> m (SymbolTable (Wise SOACS), Stms SOACS)
simplifyConsts =
  SimpleOps SOACS
-> RuleBook (Wise SOACS)
-> HoistBlockers SOACS
-> Scope SOACS
-> Stms SOACS
-> m (SymbolTable (Wise SOACS), Stms SOACS)
forall (m :: * -> *) lore.
(MonadFreshNames m, SimplifiableLore lore) =>
SimpleOps lore
-> RuleBook (Wise lore)
-> HoistBlockers lore
-> Scope lore
-> Stms lore
-> m (SymbolTable (Wise lore), Stms lore)
Simplify.simplifyStms SimpleOps SOACS
simpleSOACS RuleBook (Wise SOACS)
soacRules HoistBlockers SOACS
forall lore. HoistBlockers lore
Engine.noExtraHoistBlockers Scope SOACS
forall a. Monoid a => a
mempty

simplifySOAC :: Simplify.SimplifiableLore lore =>
                Simplify.SimplifyOp lore (SOAC lore)
simplifySOAC :: SimplifyOp lore (SOAC lore)
simplifySOAC (CmpThreshold SubExp
what String
s) = do
  SubExp
what' <- SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
what
  (SOAC (Wise lore), Stms (Wise lore))
-> SimpleM lore (SOAC (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> String -> SOAC (Wise lore)
forall lore. SubExp -> String -> SOAC lore
CmpThreshold SubExp
what' String
s, Stms (Wise lore)
forall a. Monoid a => a
mempty)
simplifySOAC (Stream SubExp
outerdim StreamForm lore
form Lambda lore
lam [VName]
arr) = do
  SubExp
outerdim' <- SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
outerdim
  (StreamForm (Wise lore)
form', Stms (Wise lore)
form_hoisted) <- StreamForm lore
-> SimpleM lore (StreamForm (Wise lore), Stms (Wise lore))
forall lore.
(Attributes lore, Simplifiable (LetAttr lore),
 Simplifiable (FParamAttr lore), Simplifiable (LParamAttr lore),
 Simplifiable (RetType lore), Simplifiable (BranchType lore),
 CanBeWise (Op lore), IndexOp (OpWithWisdom (Op lore)),
 BinderOps (Wise lore)) =>
StreamForm lore
-> SimpleM lore (StreamForm (Wise lore), Stms (Wise lore))
simplifyStreamForm StreamForm lore
form
  [VName]
arr' <- (VName -> SimpleM lore VName) -> [VName] -> SimpleM lore [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> SimpleM lore VName
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [VName]
arr
  (Lambda (Wise lore)
lam', Stms (Wise lore)
lam_hoisted) <- Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
lam ((VName -> Maybe VName) -> [VName] -> [Maybe VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Maybe VName
forall a. a -> Maybe a
Just [VName]
arr)
  (SOAC (Wise lore), Stms (Wise lore))
-> SimpleM lore (SOAC (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
-> StreamForm (Wise lore)
-> Lambda (Wise lore)
-> [VName]
-> SOAC (Wise lore)
forall lore.
SubExp -> StreamForm lore -> Lambda lore -> [VName] -> SOAC lore
Stream SubExp
outerdim' StreamForm (Wise lore)
form' Lambda (Wise lore)
lam' [VName]
arr', Stms (Wise lore)
form_hoisted Stms (Wise lore) -> Stms (Wise lore) -> Stms (Wise lore)
forall a. Semigroup a => a -> a -> a
<> Stms (Wise lore)
lam_hoisted)
  where simplifyStreamForm :: StreamForm lore
-> SimpleM lore (StreamForm (Wise lore), Stms (Wise lore))
simplifyStreamForm (Parallel StreamOrd
o Commutativity
comm Lambda lore
lam0 [SubExp]
acc) = do
          [SubExp]
acc'  <- (SubExp -> SimpleM lore SubExp)
-> [SubExp] -> SimpleM lore [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [SubExp]
acc
          (Lambda (Wise lore)
lam0', Stms (Wise lore)
hoisted) <- Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
lam0 ([Maybe VName]
 -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)))
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall a b. (a -> b) -> a -> b
$
                              Int -> Maybe VName -> [Maybe VName]
forall a. Int -> a -> [a]
replicate ([Param (LParamAttr lore)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Param (LParamAttr lore)] -> Int)
-> [Param (LParamAttr lore)] -> Int
forall a b. (a -> b) -> a -> b
$ Lambda lore -> [Param (LParamAttr lore)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
lam0) Maybe VName
forall a. Maybe a
Nothing
          (StreamForm (Wise lore), Stms (Wise lore))
-> SimpleM lore (StreamForm (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (StreamOrd
-> Commutativity
-> Lambda (Wise lore)
-> [SubExp]
-> StreamForm (Wise lore)
forall lore.
StreamOrd
-> Commutativity -> Lambda lore -> [SubExp] -> StreamForm lore
Parallel StreamOrd
o Commutativity
comm Lambda (Wise lore)
lam0' [SubExp]
acc', Stms (Wise lore)
hoisted)
        simplifyStreamForm (Sequential [SubExp]
acc) = do
          [SubExp]
acc' <- (SubExp -> SimpleM lore SubExp)
-> [SubExp] -> SimpleM lore [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [SubExp]
acc
          (StreamForm (Wise lore), Stms (Wise lore))
-> SimpleM lore (StreamForm (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp] -> StreamForm (Wise lore)
forall lore. [SubExp] -> StreamForm lore
Sequential [SubExp]
acc', Stms (Wise lore)
forall a. Monoid a => a
mempty)

simplifySOAC (Scatter SubExp
len Lambda lore
lam [VName]
ivs [(SubExp, Int, VName)]
as) = do
  SubExp
len' <- SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
len
  (Lambda (Wise lore)
lam', Stms (Wise lore)
hoisted) <- Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
lam ([Maybe VName]
 -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)))
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe VName) -> [VName] -> [Maybe VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Maybe VName
forall a. a -> Maybe a
Just [VName]
ivs
  [VName]
ivs' <- (VName -> SimpleM lore VName) -> [VName] -> SimpleM lore [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> SimpleM lore VName
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [VName]
ivs
  [(SubExp, Int, VName)]
as' <- ((SubExp, Int, VName) -> SimpleM lore (SubExp, Int, VName))
-> [(SubExp, Int, VName)] -> SimpleM lore [(SubExp, Int, VName)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (SubExp, Int, VName) -> SimpleM lore (SubExp, Int, VName)
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [(SubExp, Int, VName)]
as
  (SOAC (Wise lore), Stms (Wise lore))
-> SimpleM lore (SOAC (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
-> Lambda (Wise lore)
-> [VName]
-> [(SubExp, Int, VName)]
-> SOAC (Wise lore)
forall lore.
SubExp
-> Lambda lore -> [VName] -> [(SubExp, Int, VName)] -> SOAC lore
Scatter SubExp
len' Lambda (Wise lore)
lam' [VName]
ivs' [(SubExp, Int, VName)]
as', Stms (Wise lore)
hoisted)

simplifySOAC (Hist SubExp
w [HistOp lore]
ops Lambda lore
bfun [VName]
imgs) = do
  SubExp
w' <- SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
w
  ([HistOp (Wise lore)]
ops', [Stms (Wise lore)]
hoisted) <- ([(HistOp (Wise lore), Stms (Wise lore))]
 -> ([HistOp (Wise lore)], [Stms (Wise lore)]))
-> SimpleM lore [(HistOp (Wise lore), Stms (Wise lore))]
-> SimpleM lore ([HistOp (Wise lore)], [Stms (Wise lore)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(HistOp (Wise lore), Stms (Wise lore))]
-> ([HistOp (Wise lore)], [Stms (Wise lore)])
forall a b. [(a, b)] -> ([a], [b])
unzip (SimpleM lore [(HistOp (Wise lore), Stms (Wise lore))]
 -> SimpleM lore ([HistOp (Wise lore)], [Stms (Wise lore)]))
-> SimpleM lore [(HistOp (Wise lore), Stms (Wise lore))]
-> SimpleM lore ([HistOp (Wise lore)], [Stms (Wise lore)])
forall a b. (a -> b) -> a -> b
$ [HistOp lore]
-> (HistOp lore
    -> SimpleM lore (HistOp (Wise lore), Stms (Wise lore)))
-> SimpleM lore [(HistOp (Wise lore), Stms (Wise lore))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [HistOp lore]
ops ((HistOp lore
  -> SimpleM lore (HistOp (Wise lore), Stms (Wise lore)))
 -> SimpleM lore [(HistOp (Wise lore), Stms (Wise lore))])
-> (HistOp lore
    -> SimpleM lore (HistOp (Wise lore), Stms (Wise lore)))
-> SimpleM lore [(HistOp (Wise lore), Stms (Wise lore))]
forall a b. (a -> b) -> a -> b
$ \(HistOp SubExp
dests_w SubExp
rf [VName]
dests [SubExp]
nes Lambda lore
op) -> do
    SubExp
dests_w' <- SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
dests_w
    SubExp
rf' <- SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
rf
    [VName]
dests' <- [VName] -> SimpleM lore [VName]
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [VName]
dests
    [SubExp]
nes' <- (SubExp -> SimpleM lore SubExp)
-> [SubExp] -> SimpleM lore [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [SubExp]
nes
    (Lambda (Wise lore)
op', Stms (Wise lore)
hoisted) <- Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
op ([Maybe VName]
 -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)))
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall a b. (a -> b) -> a -> b
$ Int -> Maybe VName -> [Maybe VName]
forall a. Int -> a -> [a]
replicate ([Param (LParamAttr lore)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Param (LParamAttr lore)] -> Int)
-> [Param (LParamAttr lore)] -> Int
forall a b. (a -> b) -> a -> b
$ Lambda lore -> [Param (LParamAttr lore)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
op) Maybe VName
forall a. Maybe a
Nothing
    (HistOp (Wise lore), Stms (Wise lore))
-> SimpleM lore (HistOp (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
-> SubExp
-> [VName]
-> [SubExp]
-> Lambda (Wise lore)
-> HistOp (Wise lore)
forall lore.
SubExp
-> SubExp -> [VName] -> [SubExp] -> Lambda lore -> HistOp lore
HistOp SubExp
dests_w' SubExp
rf' [VName]
dests' [SubExp]
nes' Lambda (Wise lore)
op', Stms (Wise lore)
hoisted)
  [VName]
imgs'  <- (VName -> SimpleM lore VName) -> [VName] -> SimpleM lore [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> SimpleM lore VName
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [VName]
imgs
  (Lambda (Wise lore)
bfun', Stms (Wise lore)
bfun_hoisted) <- Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
bfun ([Maybe VName]
 -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)))
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe VName) -> [VName] -> [Maybe VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Maybe VName
forall a. a -> Maybe a
Just [VName]
imgs
  (SOAC (Wise lore), Stms (Wise lore))
-> SimpleM lore (SOAC (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
-> [HistOp (Wise lore)]
-> Lambda (Wise lore)
-> [VName]
-> SOAC (Wise lore)
forall lore.
SubExp -> [HistOp lore] -> Lambda lore -> [VName] -> SOAC lore
Hist SubExp
w' [HistOp (Wise lore)]
ops' Lambda (Wise lore)
bfun' [VName]
imgs', [Stms (Wise lore)] -> Stms (Wise lore)
forall a. Monoid a => [a] -> a
mconcat [Stms (Wise lore)]
hoisted Stms (Wise lore) -> Stms (Wise lore) -> Stms (Wise lore)
forall a. Semigroup a => a -> a -> a
<> Stms (Wise lore)
bfun_hoisted)

simplifySOAC (Screma SubExp
w (ScremaForm (Lambda lore
scan_lam, [SubExp]
scan_nes) [Reduce lore]
reds Lambda lore
map_lam) [VName]
arrs) = do
  (Lambda (Wise lore)
scan_lam', Stms (Wise lore)
scan_lam_hoisted) <-
    Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
scan_lam ([Maybe VName]
 -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)))
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall a b. (a -> b) -> a -> b
$ Int -> Maybe VName -> [Maybe VName]
forall a. Int -> a -> [a]
replicate ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
scan_nes) Maybe VName
forall a. Maybe a
Nothing
  ([Reduce (Wise lore)]
reds', [Stms (Wise lore)]
reds_hoisted) <- ([(Reduce (Wise lore), Stms (Wise lore))]
 -> ([Reduce (Wise lore)], [Stms (Wise lore)]))
-> SimpleM lore [(Reduce (Wise lore), Stms (Wise lore))]
-> SimpleM lore ([Reduce (Wise lore)], [Stms (Wise lore)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(Reduce (Wise lore), Stms (Wise lore))]
-> ([Reduce (Wise lore)], [Stms (Wise lore)])
forall a b. [(a, b)] -> ([a], [b])
unzip (SimpleM lore [(Reduce (Wise lore), Stms (Wise lore))]
 -> SimpleM lore ([Reduce (Wise lore)], [Stms (Wise lore)]))
-> SimpleM lore [(Reduce (Wise lore), Stms (Wise lore))]
-> SimpleM lore ([Reduce (Wise lore)], [Stms (Wise lore)])
forall a b. (a -> b) -> a -> b
$ [Reduce lore]
-> (Reduce lore
    -> SimpleM lore (Reduce (Wise lore), Stms (Wise lore)))
-> SimpleM lore [(Reduce (Wise lore), Stms (Wise lore))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Reduce lore]
reds ((Reduce lore
  -> SimpleM lore (Reduce (Wise lore), Stms (Wise lore)))
 -> SimpleM lore [(Reduce (Wise lore), Stms (Wise lore))])
-> (Reduce lore
    -> SimpleM lore (Reduce (Wise lore), Stms (Wise lore)))
-> SimpleM lore [(Reduce (Wise lore), Stms (Wise lore))]
forall a b. (a -> b) -> a -> b
$ \(Reduce Commutativity
comm Lambda lore
lam [SubExp]
nes) -> do
    (Lambda (Wise lore)
lam', Stms (Wise lore)
hoisted) <- Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
lam ([Maybe VName]
 -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)))
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall a b. (a -> b) -> a -> b
$ Int -> Maybe VName -> [Maybe VName]
forall a. Int -> a -> [a]
replicate ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) Maybe VName
forall a. Maybe a
Nothing
    [SubExp]
nes' <- [SubExp] -> SimpleM lore [SubExp]
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [SubExp]
nes
    (Reduce (Wise lore), Stms (Wise lore))
-> SimpleM lore (Reduce (Wise lore), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (Commutativity
-> Lambda (Wise lore) -> [SubExp] -> Reduce (Wise lore)
forall lore.
Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
Reduce Commutativity
comm Lambda (Wise lore)
lam' [SubExp]
nes', Stms (Wise lore)
hoisted)
  (Lambda (Wise lore)
map_lam', Stms (Wise lore)
map_lam_hoisted) <- Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall lore.
SimplifiableLore lore =>
Lambda lore
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
Engine.simplifyLambda Lambda lore
map_lam ([Maybe VName]
 -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)))
-> [Maybe VName]
-> SimpleM lore (Lambda (Wise lore), Stms (Wise lore))
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe VName) -> [VName] -> [Maybe VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Maybe VName
forall a. a -> Maybe a
Just [VName]
arrs
  (,) (SOAC (Wise lore)
 -> Stms (Wise lore) -> (SOAC (Wise lore), Stms (Wise lore)))
-> SimpleM lore (SOAC (Wise lore))
-> SimpleM
     lore (Stms (Wise lore) -> (SOAC (Wise lore), Stms (Wise lore)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> ScremaForm (Wise lore) -> [VName] -> SOAC (Wise lore)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma (SubExp -> ScremaForm (Wise lore) -> [VName] -> SOAC (Wise lore))
-> SimpleM lore SubExp
-> SimpleM
     lore (ScremaForm (Wise lore) -> [VName] -> SOAC (Wise lore))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
w SimpleM
  lore (ScremaForm (Wise lore) -> [VName] -> SOAC (Wise lore))
-> SimpleM lore (ScremaForm (Wise lore))
-> SimpleM lore ([VName] -> SOAC (Wise lore))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
           (Scan (Wise lore)
-> [Reduce (Wise lore)]
-> Lambda (Wise lore)
-> ScremaForm (Wise lore)
forall lore.
Scan lore -> [Reduce lore] -> Lambda lore -> ScremaForm lore
ScremaForm (Scan (Wise lore)
 -> [Reduce (Wise lore)]
 -> Lambda (Wise lore)
 -> ScremaForm (Wise lore))
-> SimpleM lore (Scan (Wise lore))
-> SimpleM
     lore
     ([Reduce (Wise lore)]
      -> Lambda (Wise lore) -> ScremaForm (Wise lore))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((,) Lambda (Wise lore)
scan_lam' ([SubExp] -> Scan (Wise lore))
-> SimpleM lore [SubExp] -> SimpleM lore (Scan (Wise lore))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExp] -> SimpleM lore [SubExp]
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [SubExp]
scan_nes) SimpleM
  lore
  ([Reduce (Wise lore)]
   -> Lambda (Wise lore) -> ScremaForm (Wise lore))
-> SimpleM lore [Reduce (Wise lore)]
-> SimpleM lore (Lambda (Wise lore) -> ScremaForm (Wise lore))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
             [Reduce (Wise lore)] -> SimpleM lore [Reduce (Wise lore)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Reduce (Wise lore)]
reds' SimpleM lore (Lambda (Wise lore) -> ScremaForm (Wise lore))
-> SimpleM lore (Lambda (Wise lore))
-> SimpleM lore (ScremaForm (Wise lore))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Lambda (Wise lore) -> SimpleM lore (Lambda (Wise lore))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda (Wise lore)
map_lam') SimpleM lore ([VName] -> SOAC (Wise lore))
-> SimpleM lore [VName] -> SimpleM lore (SOAC (Wise lore))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
            [VName] -> SimpleM lore [VName]
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify [VName]
arrs) SimpleM
  lore (Stms (Wise lore) -> (SOAC (Wise lore), Stms (Wise lore)))
-> SimpleM lore (Stms (Wise lore))
-> SimpleM lore (SOAC (Wise lore), Stms (Wise lore))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
    Stms (Wise lore) -> SimpleM lore (Stms (Wise lore))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms (Wise lore)
scan_lam_hoisted Stms (Wise lore) -> Stms (Wise lore) -> Stms (Wise lore)
forall a. Semigroup a => a -> a -> a
<> [Stms (Wise lore)] -> Stms (Wise lore)
forall a. Monoid a => [a] -> a
mconcat [Stms (Wise lore)]
reds_hoisted Stms (Wise lore) -> Stms (Wise lore) -> Stms (Wise lore)
forall a. Semigroup a => a -> a -> a
<> Stms (Wise lore)
map_lam_hoisted)

instance BinderOps (Wise SOACS) where
  mkExpAttrB :: Pattern (Wise SOACS)
-> Exp (Wise SOACS) -> m (ExpAttr (Wise SOACS))
mkExpAttrB = Pattern (Wise SOACS)
-> Exp (Wise SOACS) -> m (ExpAttr (Wise SOACS))
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m)) =>
Pattern (Lore m) -> Exp (Lore m) -> m (ExpAttr (Lore m))
bindableMkExpAttrB
  mkBodyB :: Stms (Wise SOACS) -> [SubExp] -> m (Body (Wise SOACS))
mkBodyB = Stms (Wise SOACS) -> [SubExp] -> m (Body (Wise SOACS))
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m)) =>
Stms (Lore m) -> [SubExp] -> m (Body (Lore m))
bindableMkBodyB
  mkLetNamesB :: [VName] -> Exp (Wise SOACS) -> m (Stm (Wise SOACS))
mkLetNamesB = [VName] -> Exp (Wise SOACS) -> m (Stm (Wise SOACS))
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m)) =>
[VName] -> Exp (Lore m) -> m (Stm (Lore m))
bindableMkLetNamesB

fixLambdaParams :: (MonadBinder m, Bindable (Lore m), BinderOps (Lore m)) =>
                   AST.Lambda (Lore m) -> [Maybe SubExp] -> m (AST.Lambda (Lore m))
fixLambdaParams :: Lambda (Lore m) -> [Maybe SubExp] -> m (Lambda (Lore m))
fixLambdaParams Lambda (Lore m)
lam [Maybe SubExp]
fixes = do
  Body (Lore m)
body <- Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m)))
-> Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ Scope (Lore m)
-> Binder (Lore m) (Body (Lore m))
-> Binder (Lore m) (Body (Lore m))
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope (Lore m)
forall lore attr.
(LParamAttr lore ~ attr) =>
[Param attr] -> Scope lore
scopeOfLParams ([Param Type] -> Scope (Lore m)) -> [Param Type] -> Scope (Lore m)
forall a b. (a -> b) -> a -> b
$ Lambda (Lore m) -> [LParam (Lore m)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Lore m)
lam) (Binder (Lore m) (Body (Lore m))
 -> Binder (Lore m) (Body (Lore m)))
-> Binder (Lore m) (Body (Lore m))
-> Binder (Lore m) (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ do
    (Param Type
 -> Maybe SubExp -> BinderT (Lore m) (State VNameSource) ())
-> [Param Type]
-> [Maybe SubExp]
-> BinderT (Lore m) (State VNameSource) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Param Type
-> Maybe SubExp -> BinderT (Lore m) (State VNameSource) ()
forall (m :: * -> *) attr.
MonadBinder m =>
Param attr -> Maybe SubExp -> m ()
maybeFix (Lambda (Lore m) -> [LParam (Lore m)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Lore m)
lam) [Maybe SubExp]
fixes'
    Body (Lore m) -> Binder (Lore m) (Body (Lore m))
forall (m :: * -> *) a. Monad m => a -> m a
return (Body (Lore m) -> Binder (Lore m) (Body (Lore m)))
-> Body (Lore m) -> Binder (Lore m) (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ Lambda (Lore m) -> Body (Lore m)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Lore m)
lam
  Lambda (Lore m) -> m (Lambda (Lore m))
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda (Lore m)
lam { lambdaBody :: Body (Lore m)
lambdaBody = Body (Lore m)
body
             , lambdaParams :: [LParam (Lore m)]
lambdaParams = ((LParam (Lore m), Maybe SubExp) -> LParam (Lore m))
-> [(LParam (Lore m), Maybe SubExp)] -> [LParam (Lore m)]
forall a b. (a -> b) -> [a] -> [b]
map (LParam (Lore m), Maybe SubExp) -> LParam (Lore m)
forall a b. (a, b) -> a
fst ([(LParam (Lore m), Maybe SubExp)] -> [LParam (Lore m)])
-> [(LParam (Lore m), Maybe SubExp)] -> [LParam (Lore m)]
forall a b. (a -> b) -> a -> b
$ ((LParam (Lore m), Maybe SubExp) -> Bool)
-> [(LParam (Lore m), Maybe SubExp)]
-> [(LParam (Lore m), Maybe SubExp)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Maybe SubExp -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe SubExp -> Bool)
-> ((LParam (Lore m), Maybe SubExp) -> Maybe SubExp)
-> (LParam (Lore m), Maybe SubExp)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LParam (Lore m), Maybe SubExp) -> Maybe SubExp
forall a b. (a, b) -> b
snd) ([(LParam (Lore m), Maybe SubExp)]
 -> [(LParam (Lore m), Maybe SubExp)])
-> [(LParam (Lore m), Maybe SubExp)]
-> [(LParam (Lore m), Maybe SubExp)]
forall a b. (a -> b) -> a -> b
$
                              [LParam (Lore m)]
-> [Maybe SubExp] -> [(LParam (Lore m), Maybe SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda (Lore m) -> [LParam (Lore m)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Lore m)
lam) [Maybe SubExp]
fixes' }
  where fixes' :: [Maybe SubExp]
fixes' = [Maybe SubExp]
fixes [Maybe SubExp] -> [Maybe SubExp] -> [Maybe SubExp]
forall a. [a] -> [a] -> [a]
++ Maybe SubExp -> [Maybe SubExp]
forall a. a -> [a]
repeat Maybe SubExp
forall a. Maybe a
Nothing
        maybeFix :: Param attr -> Maybe SubExp -> m ()
maybeFix Param attr
p (Just SubExp
x) = [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [Param attr -> VName
forall attr. Param attr -> VName
paramName Param attr
p] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> BasicOp (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Lore m)
forall lore. SubExp -> BasicOp lore
SubExp SubExp
x
        maybeFix Param attr
_ Maybe SubExp
Nothing = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

removeLambdaResults :: [Bool] -> AST.Lambda lore -> AST.Lambda lore
removeLambdaResults :: [Bool] -> Lambda lore -> Lambda lore
removeLambdaResults [Bool]
keep Lambda lore
lam = Lambda lore
lam { lambdaBody :: BodyT lore
lambdaBody = BodyT lore
lam_body'
                                   , lambdaReturnType :: [Type]
lambdaReturnType = [Type]
ret }
  where keep' :: [a] -> [a]
        keep' :: [a] -> [a]
keep' = ((Bool, a) -> a) -> [(Bool, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Bool, a) -> a
forall a b. (a, b) -> b
snd ([(Bool, a)] -> [a]) -> ([a] -> [(Bool, a)]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Bool, a) -> Bool) -> [(Bool, a)] -> [(Bool, a)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool, a) -> Bool
forall a b. (a, b) -> a
fst ([(Bool, a)] -> [(Bool, a)])
-> ([a] -> [(Bool, a)]) -> [a] -> [(Bool, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> [a] -> [(Bool, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Bool]
keep [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ Bool -> [Bool]
forall a. a -> [a]
repeat Bool
True)
        lam_body :: BodyT lore
lam_body = Lambda lore -> BodyT lore
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda lore
lam
        lam_body' :: BodyT lore
lam_body' = BodyT lore
lam_body { bodyResult :: [SubExp]
bodyResult = [SubExp] -> [SubExp]
forall a. [a] -> [a]
keep' ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ BodyT lore -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult BodyT lore
lam_body }
        ret :: [Type]
ret = [Type] -> [Type]
forall a. [a] -> [a]
keep' ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda lore -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda lore
lam

soacRules :: RuleBook (Wise SOACS)
soacRules :: RuleBook (Wise SOACS)
soacRules = RuleBook (Wise SOACS)
forall lore. (BinderOps lore, Aliased lore) => RuleBook lore
standardRules RuleBook (Wise SOACS)
-> RuleBook (Wise SOACS) -> RuleBook (Wise SOACS)
forall a. Semigroup a => a -> a -> a
<> [TopDownRule (Wise SOACS)]
-> [BottomUpRule (Wise SOACS)] -> RuleBook (Wise SOACS)
forall m. [TopDownRule m] -> [BottomUpRule m] -> RuleBook m
ruleBook [TopDownRule (Wise SOACS)]
topDownRules [BottomUpRule (Wise SOACS)]
bottomUpRules

topDownRules :: [TopDownRule (Wise SOACS)]
topDownRules :: [TopDownRule (Wise SOACS)]
topDownRules = [RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
hoistCertificates,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
removeReplicateMapping,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
removeReplicateWrite,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
removeUnusedSOACInput,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
simplifyClosedFormReduce,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
simplifyKnownIterationSOAC,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
fuseConcatScatter,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
simplifyMapIota,
                RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
-> TopDownRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
moveTransformToInput
               ]

bottomUpRules :: [BottomUpRule (Wise SOACS)]
bottomUpRules :: [BottomUpRule (Wise SOACS)]
bottomUpRules = [RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDeadMapping,
                 RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDeadReduction,
                 RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDeadWrite,
                 RuleBasicOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleBasicOp lore a -> SimplificationRule lore a
RuleBasicOp RuleBasicOp (Wise SOACS) (BottomUp (Wise SOACS))
forall lore. BinderOps lore => BottomUpRuleBasicOp lore
removeUnnecessaryCopy,
                 RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
liftIdentityMapping,
                 RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
liftIdentityStreaming,
                 RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDuplicateMapOutput,
                 RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
-> BottomUpRule (Wise SOACS)
forall lore a. RuleOp lore a -> SimplificationRule lore a
RuleOp RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
mapOpToOp
                ]

-- Any certificates attached to a trivial Stm in the body might as
-- well be applied to the SOAC itself.
hoistCertificates :: TopDownRuleOp (Wise SOACS)
hoistCertificates :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
hoistCertificates SymbolTable (Wise SOACS)
vtable Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
aux Op (Wise SOACS)
soac
  | (SOAC (Wise SOACS)
soac', Certificates
hoisted) <- State Certificates (SOAC (Wise SOACS))
-> Certificates -> (SOAC (Wise SOACS), Certificates)
forall s a. State s a -> s -> (a, s)
runState (SOACMapper (Wise SOACS) (Wise SOACS) (StateT Certificates Identity)
-> SOAC (Wise SOACS) -> State Certificates (SOAC (Wise SOACS))
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore)
mapSOACM SOACMapper (Wise SOACS) (Wise SOACS) (StateT Certificates Identity)
mapper Op (Wise SOACS)
SOAC (Wise SOACS)
soac) Certificates
forall a. Monoid a => a
mempty,
    Certificates
hoisted Certificates -> Certificates -> Bool
forall a. Eq a => a -> a -> Bool
/= Certificates
forall a. Monoid a => a
mempty =
      RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Certificates -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (Certificates
hoisted Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux) (RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op Op (Wise SOACS)
SOAC (Wise SOACS)
soac'
  where mapper :: SOACMapper (Wise SOACS) (Wise SOACS) (StateT Certificates Identity)
mapper = SOACMapper Any Any (StateT Certificates Identity)
forall (m :: * -> *) lore. Monad m => SOACMapper lore lore m
identitySOACMapper { mapOnSOACLambda :: Lambda (Wise SOACS)
-> StateT Certificates Identity (Lambda (Wise SOACS))
mapOnSOACLambda = Lambda (Wise SOACS)
-> StateT Certificates Identity (Lambda (Wise SOACS))
onLambda }
        onLambda :: Lambda (Wise SOACS)
-> StateT Certificates Identity (Lambda (Wise SOACS))
onLambda Lambda (Wise SOACS)
lam = do
          Stms (Wise SOACS)
stms' <- (Stm (Wise SOACS)
 -> StateT Certificates Identity (Stm (Wise SOACS)))
-> Stms (Wise SOACS)
-> StateT Certificates Identity (Stms (Wise SOACS))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Stm (Wise SOACS) -> StateT Certificates Identity (Stm (Wise SOACS))
onStm (Stms (Wise SOACS)
 -> StateT Certificates Identity (Stms (Wise SOACS)))
-> Stms (Wise SOACS)
-> StateT Certificates Identity (Stms (Wise SOACS))
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> Stms (Wise SOACS)
forall lore. BodyT lore -> Stms lore
bodyStms (Body (Wise SOACS) -> Stms (Wise SOACS))
-> Body (Wise SOACS) -> Stms (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
lam
          Lambda (Wise SOACS)
-> StateT Certificates Identity (Lambda (Wise SOACS))
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda (Wise SOACS)
lam { lambdaBody :: Body (Wise SOACS)
lambdaBody =
                       Stms (Wise SOACS) -> [SubExp] -> Body (Wise SOACS)
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms (Wise SOACS)
stms' ([SubExp] -> Body (Wise SOACS)) -> [SubExp] -> Body (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
lam }
        onStm :: Stm (Wise SOACS) -> StateT Certificates Identity (Stm (Wise SOACS))
onStm (Let Pattern (Wise SOACS)
se_pat StmAux (ExpAttr (Wise SOACS))
se_aux (BasicOp (SubExp SubExp
se))) = do
          let ([VName]
invariant, [VName]
variant) =
                (VName -> Bool) -> [VName] -> ([VName], [VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (VName -> SymbolTable (Wise SOACS) -> Bool
forall lore. VName -> SymbolTable lore -> Bool
`ST.elem` SymbolTable (Wise SOACS)
vtable) ([VName] -> ([VName], [VName])) -> [VName] -> ([VName], [VName])
forall a b. (a -> b) -> a -> b
$
                Certificates -> [VName]
unCertificates (Certificates -> [VName]) -> Certificates -> [VName]
forall a b. (a -> b) -> a -> b
$ StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
se_aux
              se_aux' :: StmAux (ExpWisdom, ())
se_aux' = StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
se_aux { stmAuxCerts :: Certificates
stmAuxCerts = [VName] -> Certificates
Certificates [VName]
variant }
          (Certificates -> Certificates) -> StateT Certificates Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ([VName] -> Certificates
Certificates [VName]
invariantCertificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<>)
          Stm (Wise SOACS) -> StateT Certificates Identity (Stm (Wise SOACS))
forall (m :: * -> *) a. Monad m => a -> m a
return (Stm (Wise SOACS)
 -> StateT Certificates Identity (Stm (Wise SOACS)))
-> Stm (Wise SOACS)
-> StateT Certificates Identity (Stm (Wise SOACS))
forall a b. (a -> b) -> a -> b
$ Pattern (Wise SOACS)
-> StmAux (ExpAttr (Wise SOACS))
-> Exp (Wise SOACS)
-> Stm (Wise SOACS)
forall lore.
Pattern lore -> StmAux (ExpAttr lore) -> Exp lore -> Stm lore
Let Pattern (Wise SOACS)
se_pat StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
se_aux' (Exp (Wise SOACS) -> Stm (Wise SOACS))
-> Exp (Wise SOACS) -> Stm (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Wise SOACS)
forall lore. SubExp -> BasicOp lore
SubExp SubExp
se
        onStm Stm (Wise SOACS)
stm = Stm (Wise SOACS) -> StateT Certificates Identity (Stm (Wise SOACS))
forall (m :: * -> *) a. Monad m => a -> m a
return Stm (Wise SOACS)
stm
hoistCertificates SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ =
  Rule (Wise SOACS)
forall lore. Rule lore
Skip

liftIdentityMapping :: BottomUpRuleOp (Wise SOACS)
liftIdentityMapping :: RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
liftIdentityMapping (SymbolTable (Wise SOACS)
_, UsageTable
usages) Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma w form arrs)
  | Just Lambda (Wise SOACS)
fun <- ScremaForm (Wise SOACS) -> Maybe (Lambda (Wise SOACS))
forall lore. ScremaForm lore -> Maybe (Lambda lore)
isMapSOAC ScremaForm (Wise SOACS)
form = do
  let inputMap :: Map VName VName
inputMap = [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$ [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall attr. Param attr -> VName
paramName ([Param Type] -> [VName]) -> [Param Type] -> [VName]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
fun) [VName]
arrs
      free :: Names
free = Body (Wise SOACS) -> Names
forall a. FreeIn a => a -> Names
freeIn (Body (Wise SOACS) -> Names) -> Body (Wise SOACS) -> Names
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun
      rettype :: [Type]
rettype = Lambda (Wise SOACS) -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda (Wise SOACS)
fun
      ses :: [SubExp]
ses = Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun

      freeOrConst :: SubExp -> Bool
freeOrConst (Var VName
v)    = VName
v VName -> Names -> Bool
`nameIn` Names
free
      freeOrConst Constant{} = Bool
True

      checkInvariance :: (PatElemT (VarWisdom, Type), SubExp, Type)
-> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
    [(PatElemT (VarWisdom, Type), SubExp)], [Type])
-> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
    [(PatElemT (VarWisdom, Type), SubExp)], [Type])
checkInvariance (PatElemT (VarWisdom, Type)
outId, Var VName
v, Type
_) ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
invariant, [(PatElemT (VarWisdom, Type), SubExp)]
mapresult, [Type]
rettype')
        | Just VName
inp <- VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
inputMap =
            let e :: BasicOp (Wise SOACS)
e | PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
outId VName -> UsageTable -> Bool
`UT.isConsumed` UsageTable
usages
                    Bool -> Bool -> Bool
|| VName
inp VName -> UsageTable -> Bool
`UT.isConsumed` UsageTable
usages =
                      VName -> BasicOp (Wise SOACS)
forall lore. VName -> BasicOp lore
Copy VName
inp
                  | Bool
otherwise =
                      SubExp -> BasicOp (Wise SOACS)
forall lore. SubExp -> BasicOp lore
SubExp (SubExp -> BasicOp (Wise SOACS)) -> SubExp -> BasicOp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
inp
            in (([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)
outId], BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp BasicOp (Wise SOACS)
e) (PatternT (VarWisdom, Type), Exp (Wise SOACS))
-> [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
-> [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
forall a. a -> [a] -> [a]
: [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
invariant,
                [(PatElemT (VarWisdom, Type), SubExp)]
mapresult,
                [Type]
rettype')
      checkInvariance (PatElemT (VarWisdom, Type)
outId, SubExp
e, Type
t) ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
invariant, [(PatElemT (VarWisdom, Type), SubExp)]
mapresult, [Type]
rettype')
        | SubExp -> Bool
freeOrConst SubExp
e = (([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)
outId], BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp (Wise SOACS)
forall lore. Shape -> SubExp -> BasicOp lore
Replicate ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) SubExp
e) (PatternT (VarWisdom, Type), Exp (Wise SOACS))
-> [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
-> [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
forall a. a -> [a] -> [a]
: [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
invariant,
                           [(PatElemT (VarWisdom, Type), SubExp)]
mapresult,
                           [Type]
rettype')
        | Bool
otherwise = ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
invariant,
                       (PatElemT (VarWisdom, Type)
outId, SubExp
e) (PatElemT (VarWisdom, Type), SubExp)
-> [(PatElemT (VarWisdom, Type), SubExp)]
-> [(PatElemT (VarWisdom, Type), SubExp)]
forall a. a -> [a] -> [a]
: [(PatElemT (VarWisdom, Type), SubExp)]
mapresult,
                       Type
t Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
rettype')

  case ((PatElemT (VarWisdom, Type), SubExp, Type)
 -> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
     [(PatElemT (VarWisdom, Type), SubExp)], [Type])
 -> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
     [(PatElemT (VarWisdom, Type), SubExp)], [Type]))
-> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
    [(PatElemT (VarWisdom, Type), SubExp)], [Type])
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
-> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
    [(PatElemT (VarWisdom, Type), SubExp)], [Type])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (PatElemT (VarWisdom, Type), SubExp, Type)
-> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
    [(PatElemT (VarWisdom, Type), SubExp)], [Type])
-> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
    [(PatElemT (VarWisdom, Type), SubExp)], [Type])
checkInvariance ([], [], []) ([(PatElemT (VarWisdom, Type), SubExp, Type)]
 -> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
     [(PatElemT (VarWisdom, Type), SubExp)], [Type]))
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
-> ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))],
    [(PatElemT (VarWisdom, Type), SubExp)], [Type])
forall a b. (a -> b) -> a -> b
$
       [PatElemT (VarWisdom, Type)]
-> [SubExp]
-> [Type]
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternElements PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat) [SubExp]
ses [Type]
rettype of
    ([], [(PatElemT (VarWisdom, Type), SubExp)]
_, [Type]
_) -> Rule (Wise SOACS)
forall lore. Rule lore
Skip
    ([(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
invariant, [(PatElemT (VarWisdom, Type), SubExp)]
mapresult, [Type]
rettype') -> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      let ([PatElemT (VarWisdom, Type)]
pat', [SubExp]
ses') = [(PatElemT (VarWisdom, Type), SubExp)]
-> ([PatElemT (VarWisdom, Type)], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(PatElemT (VarWisdom, Type), SubExp)]
mapresult
          fun' :: Lambda (Wise SOACS)
fun' = Lambda (Wise SOACS)
fun { lambdaBody :: Body (Wise SOACS)
lambdaBody = (Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun) { bodyResult :: [SubExp]
bodyResult = [SubExp]
ses' }
                     , lambdaReturnType :: [Type]
lambdaReturnType = [Type]
rettype'
                     }
      ((PatternT (VarWisdom, Type), Exp (Wise SOACS))
 -> RuleM (Wise SOACS) [Ident])
-> [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
-> RuleM (Wise SOACS) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((PatternT (VarWisdom, Type)
 -> Exp (Wise SOACS) -> RuleM (Wise SOACS) [Ident])
-> (PatternT (VarWisdom, Type), Exp (Wise SOACS))
-> RuleM (Wise SOACS) [Ident]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry PatternT (VarWisdom, Type)
-> Exp (Wise SOACS) -> RuleM (Wise SOACS) [Ident]
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m [Ident]
letBind) [(PatternT (VarWisdom, Type), Exp (Wise SOACS))]
invariant
      [VName] -> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ ((PatElemT (VarWisdom, Type) -> VName)
-> [PatElemT (VarWisdom, Type)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName [PatElemT (VarWisdom, Type)]
pat') (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w (Lambda (Wise SOACS) -> ScremaForm (Wise SOACS)
forall lore. Bindable lore => Lambda lore -> ScremaForm lore
mapSOAC Lambda (Wise SOACS)
fun') [VName]
arrs
liftIdentityMapping BottomUp (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS)
liftIdentityStreaming :: RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
liftIdentityStreaming BottomUp (Wise SOACS)
_ (Pattern [] [PatElemT (LetAttr (Wise SOACS))]
pes) StmAux (ExpAttr (Wise SOACS))
_ (Stream w form lam arrs)
  | ([(Type, PatElemT (VarWisdom, Type), SubExp)]
variant_map, [(PatElemT (VarWisdom, Type), VName)]
invariant_map) <-
      [Either
   (Type, PatElemT (VarWisdom, Type), SubExp)
   (PatElemT (VarWisdom, Type), VName)]
-> ([(Type, PatElemT (VarWisdom, Type), SubExp)],
    [(PatElemT (VarWisdom, Type), VName)])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either
    (Type, PatElemT (VarWisdom, Type), SubExp)
    (PatElemT (VarWisdom, Type), VName)]
 -> ([(Type, PatElemT (VarWisdom, Type), SubExp)],
     [(PatElemT (VarWisdom, Type), VName)]))
-> [Either
      (Type, PatElemT (VarWisdom, Type), SubExp)
      (PatElemT (VarWisdom, Type), VName)]
-> ([(Type, PatElemT (VarWisdom, Type), SubExp)],
    [(PatElemT (VarWisdom, Type), VName)])
forall a b. (a -> b) -> a -> b
$ ((Type, PatElemT (VarWisdom, Type), SubExp)
 -> Either
      (Type, PatElemT (VarWisdom, Type), SubExp)
      (PatElemT (VarWisdom, Type), VName))
-> [(Type, PatElemT (VarWisdom, Type), SubExp)]
-> [Either
      (Type, PatElemT (VarWisdom, Type), SubExp)
      (PatElemT (VarWisdom, Type), VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Type, PatElemT (VarWisdom, Type), SubExp)
-> Either
     (Type, PatElemT (VarWisdom, Type), SubExp)
     (PatElemT (VarWisdom, Type), VName)
isInvariantRes ([(Type, PatElemT (VarWisdom, Type), SubExp)]
 -> [Either
       (Type, PatElemT (VarWisdom, Type), SubExp)
       (PatElemT (VarWisdom, Type), VName)])
-> [(Type, PatElemT (VarWisdom, Type), SubExp)]
-> [Either
      (Type, PatElemT (VarWisdom, Type), SubExp)
      (PatElemT (VarWisdom, Type), VName)]
forall a b. (a -> b) -> a -> b
$ [Type]
-> [PatElemT (VarWisdom, Type)]
-> [SubExp]
-> [(Type, PatElemT (VarWisdom, Type), SubExp)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Type]
map_ts [PatElemT (VarWisdom, Type)]
map_pes [SubExp]
map_res,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [(PatElemT (VarWisdom, Type), VName)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(PatElemT (VarWisdom, Type), VName)]
invariant_map = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do

      [(PatElemT (VarWisdom, Type), VName)]
-> ((PatElemT (VarWisdom, Type), VName) -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(PatElemT (VarWisdom, Type), VName)]
invariant_map (((PatElemT (VarWisdom, Type), VName) -> RuleM (Wise SOACS) ())
 -> RuleM (Wise SOACS) ())
-> ((PatElemT (VarWisdom, Type), VName) -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ \(PatElemT (VarWisdom, Type)
pe, VName
arr) ->
        Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ ([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)
pe]) (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp (Wise SOACS)
forall lore. VName -> BasicOp lore
Copy VName
arr

      let ([Type]
variant_map_ts, [PatElemT (VarWisdom, Type)]
variant_map_pes, [SubExp]
variant_map_res) = [(Type, PatElemT (VarWisdom, Type), SubExp)]
-> ([Type], [PatElemT (VarWisdom, Type)], [SubExp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(Type, PatElemT (VarWisdom, Type), SubExp)]
variant_map
          lam' :: Lambda (Wise SOACS)
lam' = Lambda (Wise SOACS)
lam { lambdaBody :: Body (Wise SOACS)
lambdaBody = (Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
lam) { bodyResult :: [SubExp]
bodyResult = [SubExp]
fold_res [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
variant_map_res }
                     , lambdaReturnType :: [Type]
lambdaReturnType = [Type]
fold_ts [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
variant_map_ts }

      Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ ([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] ([PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type))
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall a b. (a -> b) -> a -> b
$ [PatElemT (VarWisdom, Type)]
fold_pes [PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> [PatElemT (VarWisdom, Type)]
forall a. [a] -> [a] -> [a]
++ [PatElemT (VarWisdom, Type)]
variant_map_pes) (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$
        Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp
-> StreamForm (Wise SOACS)
-> Lambda (Wise SOACS)
-> [VName]
-> SOAC (Wise SOACS)
forall lore.
SubExp -> StreamForm lore -> Lambda lore -> [VName] -> SOAC lore
Stream SubExp
w StreamForm (Wise SOACS)
form Lambda (Wise SOACS)
lam' [VName]
arrs
  where num_folds :: Int
num_folds = [SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int) -> [SubExp] -> Int
forall a b. (a -> b) -> a -> b
$ StreamForm (Wise SOACS) -> [SubExp]
forall lore. StreamForm lore -> [SubExp]
getStreamAccums StreamForm (Wise SOACS)
form
        ([PatElemT (VarWisdom, Type)]
fold_pes, [PatElemT (VarWisdom, Type)]
map_pes) = Int
-> [PatElemT (VarWisdom, Type)]
-> ([PatElemT (VarWisdom, Type)], [PatElemT (VarWisdom, Type)])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
num_folds [PatElemT (VarWisdom, Type)]
[PatElemT (LetAttr (Wise SOACS))]
pes
        ([Type]
fold_ts, [Type]
map_ts) = Int -> [Type] -> ([Type], [Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
num_folds ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda (Wise SOACS)
lam
        lam_res :: [SubExp]
lam_res = Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
lam
        ([SubExp]
fold_res, [SubExp]
map_res) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
num_folds [SubExp]
lam_res
        params_to_arrs :: [(VName, VName)]
params_to_arrs = [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall attr. Param attr -> VName
paramName ([Param Type] -> [VName]) -> [Param Type] -> [VName]
forall a b. (a -> b) -> a -> b
$ Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
drop (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
num_folds) ([Param Type] -> [Param Type]) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
lam) [VName]
arrs

        isInvariantRes :: (Type, PatElemT (VarWisdom, Type), SubExp)
-> Either
     (Type, PatElemT (VarWisdom, Type), SubExp)
     (PatElemT (VarWisdom, Type), VName)
isInvariantRes (Type
_, PatElemT (VarWisdom, Type)
pe, Var VName
v)
          | Just VName
arr <- VName -> [(VName, VName)] -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v [(VName, VName)]
params_to_arrs =
              (PatElemT (VarWisdom, Type), VName)
-> Either
     (Type, PatElemT (VarWisdom, Type), SubExp)
     (PatElemT (VarWisdom, Type), VName)
forall a b. b -> Either a b
Right (PatElemT (VarWisdom, Type)
pe, VName
arr)
        isInvariantRes (Type, PatElemT (VarWisdom, Type), SubExp)
x =
          (Type, PatElemT (VarWisdom, Type), SubExp)
-> Either
     (Type, PatElemT (VarWisdom, Type), SubExp)
     (PatElemT (VarWisdom, Type), VName)
forall a b. a -> Either a b
Left (Type, PatElemT (VarWisdom, Type), SubExp)
x
liftIdentityStreaming BottomUp (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

-- | Remove all arguments to the map that are simply replicates.
-- These can be turned into free variables instead.
removeReplicateMapping :: TopDownRuleOp (Wise SOACS)
removeReplicateMapping :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
removeReplicateMapping SymbolTable (Wise SOACS)
vtable Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma w form arrs)
  | Just Lambda (Wise SOACS)
fun <- ScremaForm (Wise SOACS) -> Maybe (Lambda (Wise SOACS))
forall lore. ScremaForm lore -> Maybe (Lambda lore)
isMapSOAC ScremaForm (Wise SOACS)
form,
    Just ([([VName], Certificates, Exp (Wise SOACS))]
bnds, Lambda (Wise SOACS)
fun', [VName]
arrs') <- SymbolTable (Wise SOACS)
-> Lambda (Wise SOACS)
-> [VName]
-> Maybe
     ([([VName], Certificates, Exp (Wise SOACS))], Lambda (Wise SOACS),
      [VName])
forall lore.
Aliased lore =>
SymbolTable lore
-> Lambda lore
-> [VName]
-> Maybe
     ([([VName], Certificates, Exp lore)], Lambda lore, [VName])
removeReplicateInput SymbolTable (Wise SOACS)
vtable Lambda (Wise SOACS)
fun [VName]
arrs = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      [([VName], Certificates, Exp (Wise SOACS))]
-> (([VName], Certificates, Exp (Wise SOACS))
    -> RuleM (Wise SOACS) [Ident])
-> RuleM (Wise SOACS) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [([VName], Certificates, Exp (Wise SOACS))]
bnds ((([VName], Certificates, Exp (Wise SOACS))
  -> RuleM (Wise SOACS) [Ident])
 -> RuleM (Wise SOACS) ())
-> (([VName], Certificates, Exp (Wise SOACS))
    -> RuleM (Wise SOACS) [Ident])
-> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ \([VName]
vs,Certificates
cs,Exp (Wise SOACS)
e) -> Certificates
-> RuleM (Wise SOACS) [Ident] -> RuleM (Wise SOACS) [Ident]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (RuleM (Wise SOACS) [Ident] -> RuleM (Wise SOACS) [Ident])
-> RuleM (Wise SOACS) [Ident] -> RuleM (Wise SOACS) [Ident]
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) [Ident]
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m [Ident]
letBindNames [VName]
vs Exp (Lore (RuleM (Wise SOACS)))
Exp (Wise SOACS)
e
      Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w (Lambda (Wise SOACS) -> ScremaForm (Wise SOACS)
forall lore. Bindable lore => Lambda lore -> ScremaForm lore
mapSOAC Lambda (Wise SOACS)
fun') [VName]
arrs'
removeReplicateMapping SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

-- | Like 'removeReplicateMapping', but for 'Scatter'.
removeReplicateWrite :: TopDownRuleOp (Wise SOACS)
removeReplicateWrite :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
removeReplicateWrite SymbolTable (Wise SOACS)
vtable Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Scatter len lam ivs as)
  | Just ([([VName], Certificates, Exp (Wise SOACS))]
bnds, Lambda (Wise SOACS)
lam', [VName]
ivs') <- SymbolTable (Wise SOACS)
-> Lambda (Wise SOACS)
-> [VName]
-> Maybe
     ([([VName], Certificates, Exp (Wise SOACS))], Lambda (Wise SOACS),
      [VName])
forall lore.
Aliased lore =>
SymbolTable lore
-> Lambda lore
-> [VName]
-> Maybe
     ([([VName], Certificates, Exp lore)], Lambda lore, [VName])
removeReplicateInput SymbolTable (Wise SOACS)
vtable Lambda (Wise SOACS)
lam [VName]
ivs = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      [([VName], Certificates, Exp (Wise SOACS))]
-> (([VName], Certificates, Exp (Wise SOACS))
    -> RuleM (Wise SOACS) [Ident])
-> RuleM (Wise SOACS) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [([VName], Certificates, Exp (Wise SOACS))]
bnds ((([VName], Certificates, Exp (Wise SOACS))
  -> RuleM (Wise SOACS) [Ident])
 -> RuleM (Wise SOACS) ())
-> (([VName], Certificates, Exp (Wise SOACS))
    -> RuleM (Wise SOACS) [Ident])
-> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ \([VName]
vs,Certificates
cs,Exp (Wise SOACS)
e) -> Certificates
-> RuleM (Wise SOACS) [Ident] -> RuleM (Wise SOACS) [Ident]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (RuleM (Wise SOACS) [Ident] -> RuleM (Wise SOACS) [Ident])
-> RuleM (Wise SOACS) [Ident] -> RuleM (Wise SOACS) [Ident]
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) [Ident]
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m [Ident]
letBindNames [VName]
vs Exp (Lore (RuleM (Wise SOACS)))
Exp (Wise SOACS)
e
      Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp
-> Lambda (Wise SOACS)
-> [VName]
-> [(SubExp, Int, VName)]
-> SOAC (Wise SOACS)
forall lore.
SubExp
-> Lambda lore -> [VName] -> [(SubExp, Int, VName)] -> SOAC lore
Scatter SubExp
len Lambda (Wise SOACS)
lam' [VName]
ivs' [(SubExp, Int, VName)]
as
removeReplicateWrite SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

removeReplicateInput :: Aliased lore =>
                        ST.SymbolTable lore
                     -> AST.Lambda lore -> [VName]
                     -> Maybe ([([VName], Certificates, AST.Exp lore)],
                                AST.Lambda lore, [VName])
removeReplicateInput :: SymbolTable lore
-> Lambda lore
-> [VName]
-> Maybe
     ([([VName], Certificates, Exp lore)], Lambda lore, [VName])
removeReplicateInput SymbolTable lore
vtable Lambda lore
fun [VName]
arrs
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [([VName], Certificates, Exp lore)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [([VName], Certificates, Exp lore)]
parameterBnds = do
  let ([Param (LParamAttr lore)]
arr_params', [VName]
arrs') = [(Param (LParamAttr lore), VName)]
-> ([Param (LParamAttr lore)], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param (LParamAttr lore), VName)]
params_and_arrs
      fun' :: Lambda lore
fun' = Lambda lore
fun { lambdaParams :: [Param (LParamAttr lore)]
lambdaParams = [Param (LParamAttr lore)]
acc_params [Param (LParamAttr lore)]
-> [Param (LParamAttr lore)] -> [Param (LParamAttr lore)]
forall a. Semigroup a => a -> a -> a
<> [Param (LParamAttr lore)]
arr_params' }
  ([([VName], Certificates, Exp lore)], Lambda lore, [VName])
-> Maybe
     ([([VName], Certificates, Exp lore)], Lambda lore, [VName])
forall (m :: * -> *) a. Monad m => a -> m a
return ([([VName], Certificates, Exp lore)]
parameterBnds, Lambda lore
fun', [VName]
arrs')
  | Bool
otherwise = Maybe ([([VName], Certificates, Exp lore)], Lambda lore, [VName])
forall a. Maybe a
Nothing

  where params :: [Param (LParamAttr lore)]
params = Lambda lore -> [Param (LParamAttr lore)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
fun
        ([Param (LParamAttr lore)]
acc_params, [Param (LParamAttr lore)]
arr_params) =
          Int
-> [Param (LParamAttr lore)]
-> ([Param (LParamAttr lore)], [Param (LParamAttr lore)])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Param (LParamAttr lore)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param (LParamAttr lore)]
params Int -> Int -> Int
forall a. Num a => a -> a -> a
- [VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
arrs) [Param (LParamAttr lore)]
params
        ([(Param (LParamAttr lore), VName)]
params_and_arrs, [([VName], Certificates, Exp lore)]
parameterBnds) =
          [Either
   (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)]
-> ([(Param (LParamAttr lore), VName)],
    [([VName], Certificates, Exp lore)])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either
    (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)]
 -> ([(Param (LParamAttr lore), VName)],
     [([VName], Certificates, Exp lore)]))
-> [Either
      (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)]
-> ([(Param (LParamAttr lore), VName)],
    [([VName], Certificates, Exp lore)])
forall a b. (a -> b) -> a -> b
$ (Param (LParamAttr lore)
 -> VName
 -> Either
      (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore))
-> [Param (LParamAttr lore)]
-> [VName]
-> [Either
      (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Param (LParamAttr lore)
-> VName
-> Either
     (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)
isReplicateAndNotConsumed [Param (LParamAttr lore)]
arr_params [VName]
arrs

        isReplicateAndNotConsumed :: Param (LParamAttr lore)
-> VName
-> Either
     (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)
isReplicateAndNotConsumed Param (LParamAttr lore)
p VName
v
          | Just (BasicOp (Replicate (Shape (SubExp
_:[SubExp]
ds)) SubExp
e), Certificates
v_cs) <-
              VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
forall lore.
VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
ST.lookupExp VName
v SymbolTable lore
vtable,
            Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Param (LParamAttr lore) -> VName
forall attr. Param attr -> VName
paramName Param (LParamAttr lore)
p VName -> Names -> Bool
`nameIn` Lambda lore -> Names
forall lore. Aliased lore => Lambda lore -> Names
consumedByLambda Lambda lore
fun =
              ([VName], Certificates, Exp lore)
-> Either
     (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)
forall a b. b -> Either a b
Right ([Param (LParamAttr lore) -> VName
forall attr. Param attr -> VName
paramName Param (LParamAttr lore)
p],
                     Certificates
v_cs,
                     case [SubExp]
ds of
                       [] -> BasicOp lore -> Exp lore
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp lore -> Exp lore) -> BasicOp lore -> Exp lore
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp lore
forall lore. SubExp -> BasicOp lore
SubExp SubExp
e
                       [SubExp]
_  -> BasicOp lore -> Exp lore
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp lore -> Exp lore) -> BasicOp lore -> Exp lore
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp lore
forall lore. Shape -> SubExp -> BasicOp lore
Replicate ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
ds) SubExp
e)
          | Bool
otherwise =
              (Param (LParamAttr lore), VName)
-> Either
     (Param (LParamAttr lore), VName) ([VName], Certificates, Exp lore)
forall a b. a -> Either a b
Left (Param (LParamAttr lore)
p, VName
v)

-- | Remove inputs that are not used inside the SOAC.
removeUnusedSOACInput :: TopDownRuleOp (Wise SOACS)
removeUnusedSOACInput :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
removeUnusedSOACInput SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma w (ScremaForm scan reduce map_lam) arrs)
  | ([(Param Type, VName)]
used,[(Param Type, VName)]
unused) <- ((Param Type, VName) -> Bool)
-> [(Param Type, VName)]
-> ([(Param Type, VName)], [(Param Type, VName)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (Param Type, VName) -> Bool
usedInput [(Param Type, VName)]
params_and_arrs,
    Bool -> Bool
not ([(Param Type, VName)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Param Type, VName)]
unused) = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      let ([Param Type]
used_params, [VName]
used_arrs) = [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param Type, VName)]
used
          map_lam' :: Lambda (Wise SOACS)
map_lam' = Lambda (Wise SOACS)
map_lam { lambdaParams :: [LParam (Wise SOACS)]
lambdaParams = [Param Type]
[LParam (Wise SOACS)]
used_params }
      Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w (Scan (Wise SOACS)
-> [Reduce (Wise SOACS)]
-> Lambda (Wise SOACS)
-> ScremaForm (Wise SOACS)
forall lore.
Scan lore -> [Reduce lore] -> Lambda lore -> ScremaForm lore
ScremaForm Scan (Wise SOACS)
scan [Reduce (Wise SOACS)]
reduce Lambda (Wise SOACS)
map_lam') [VName]
used_arrs
  where params_and_arrs :: [(Param Type, VName)]
params_and_arrs = [Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
map_lam) [VName]
arrs
        used_in_body :: Names
used_in_body = Body (Wise SOACS) -> Names
forall a. FreeIn a => a -> Names
freeIn (Body (Wise SOACS) -> Names) -> Body (Wise SOACS) -> Names
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam
        usedInput :: (Param Type, VName) -> Bool
usedInput (Param Type
param, VName
_) = Param Type -> VName
forall attr. Param attr -> VName
paramName Param Type
param VName -> Names -> Bool
`nameIn` Names
used_in_body
removeUnusedSOACInput SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

removeDeadMapping :: BottomUpRuleOp (Wise SOACS)
removeDeadMapping :: RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDeadMapping (SymbolTable (Wise SOACS)
_, UsageTable
used) Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma w form arrs)
  | Just Lambda (Wise SOACS)
fun <- ScremaForm (Wise SOACS) -> Maybe (Lambda (Wise SOACS))
forall lore. ScremaForm lore -> Maybe (Lambda lore)
isMapSOAC ScremaForm (Wise SOACS)
form =
  let ses :: [SubExp]
ses = Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun
      isUsed :: (PatElemT (VarWisdom, Type), SubExp, Type) -> Bool
isUsed (PatElemT (VarWisdom, Type)
bindee, SubExp
_, Type
_) = (VName -> UsageTable -> Bool
`UT.used` UsageTable
used) (VName -> Bool) -> VName -> Bool
forall a b. (a -> b) -> a -> b
$ PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
bindee
      ([PatElemT (VarWisdom, Type)]
pat',[SubExp]
ses', [Type]
ts') = [(PatElemT (VarWisdom, Type), SubExp, Type)]
-> ([PatElemT (VarWisdom, Type)], [SubExp], [Type])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(PatElemT (VarWisdom, Type), SubExp, Type)]
 -> ([PatElemT (VarWisdom, Type)], [SubExp], [Type]))
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
-> ([PatElemT (VarWisdom, Type)], [SubExp], [Type])
forall a b. (a -> b) -> a -> b
$ ((PatElemT (VarWisdom, Type), SubExp, Type) -> Bool)
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
forall a. (a -> Bool) -> [a] -> [a]
filter (PatElemT (VarWisdom, Type), SubExp, Type) -> Bool
isUsed ([(PatElemT (VarWisdom, Type), SubExp, Type)]
 -> [(PatElemT (VarWisdom, Type), SubExp, Type)])
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
forall a b. (a -> b) -> a -> b
$
                         [PatElemT (VarWisdom, Type)]
-> [SubExp]
-> [Type]
-> [(PatElemT (VarWisdom, Type), SubExp, Type)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternElements PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat) [SubExp]
ses ([Type] -> [(PatElemT (VarWisdom, Type), SubExp, Type)])
-> [Type] -> [(PatElemT (VarWisdom, Type), SubExp, Type)]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda (Wise SOACS)
fun
      fun' :: Lambda (Wise SOACS)
fun' = Lambda (Wise SOACS)
fun { lambdaBody :: Body (Wise SOACS)
lambdaBody = (Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun) { bodyResult :: [SubExp]
bodyResult = [SubExp]
ses' }
                 , lambdaReturnType :: [Type]
lambdaReturnType = [Type]
ts'
                 }
  in if PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat PatternT (VarWisdom, Type) -> PatternT (VarWisdom, Type) -> Bool
forall a. Eq a => a -> a -> Bool
/= [PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)]
pat'
     then RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ ([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)]
pat') (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w (Lambda (Wise SOACS) -> ScremaForm (Wise SOACS)
forall lore. Bindable lore => Lambda lore -> ScremaForm lore
mapSOAC Lambda (Wise SOACS)
fun') [VName]
arrs
     else Rule (Wise SOACS)
forall lore. Rule lore
Skip
removeDeadMapping BottomUp (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

removeDuplicateMapOutput :: BottomUpRuleOp (Wise SOACS)
removeDuplicateMapOutput :: RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDuplicateMapOutput (SymbolTable (Wise SOACS)
_, UsageTable
used) Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma w form arrs)
  | Just Lambda (Wise SOACS)
fun <- ScremaForm (Wise SOACS) -> Maybe (Lambda (Wise SOACS))
forall lore. ScremaForm lore -> Maybe (Lambda lore)
isMapSOAC ScremaForm (Wise SOACS)
form =
  let ses :: [SubExp]
ses = Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun
      ts :: [Type]
ts = Lambda (Wise SOACS) -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda (Wise SOACS)
fun
      pes :: [PatElemT (VarWisdom, Type)]
pes = PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternValueElements PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat
      ses_ts_pes :: [(SubExp, Type, PatElemT (VarWisdom, Type))]
ses_ts_pes = [SubExp]
-> [Type]
-> [PatElemT (VarWisdom, Type)]
-> [(SubExp, Type, PatElemT (VarWisdom, Type))]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [SubExp]
ses [Type]
ts [PatElemT (VarWisdom, Type)]
pes
      ([(SubExp, Type, PatElemT (VarWisdom, Type))]
ses_ts_pes', [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))]
copies) =
        (([(SubExp, Type, PatElemT (VarWisdom, Type))],
  [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))])
 -> (SubExp, Type, PatElemT (VarWisdom, Type))
 -> ([(SubExp, Type, PatElemT (VarWisdom, Type))],
     [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))]))
-> ([(SubExp, Type, PatElemT (VarWisdom, Type))],
    [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))])
-> [(SubExp, Type, PatElemT (VarWisdom, Type))]
-> ([(SubExp, Type, PatElemT (VarWisdom, Type))],
    [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))])
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ([(SubExp, Type, PatElemT (VarWisdom, Type))],
 [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))])
-> (SubExp, Type, PatElemT (VarWisdom, Type))
-> ([(SubExp, Type, PatElemT (VarWisdom, Type))],
    [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))])
forall a b a.
Eq a =>
([(a, b, a)], [(a, a)]) -> (a, b, a) -> ([(a, b, a)], [(a, a)])
checkForDuplicates ([(SubExp, Type, PatElemT (VarWisdom, Type))]
forall a. Monoid a => a
mempty,[(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))]
forall a. Monoid a => a
mempty) [(SubExp, Type, PatElemT (VarWisdom, Type))]
ses_ts_pes
  in if [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))]
copies then Rule (Wise SOACS)
forall lore. Rule lore
Skip
     else RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
       let ([SubExp]
ses', [Type]
ts', [PatElemT (VarWisdom, Type)]
pes') = [(SubExp, Type, PatElemT (VarWisdom, Type))]
-> ([SubExp], [Type], [PatElemT (VarWisdom, Type)])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(SubExp, Type, PatElemT (VarWisdom, Type))]
ses_ts_pes'
           pat' :: PatternT (VarWisdom, Type)
pat' = [PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)]
pes'
           fun' :: Lambda (Wise SOACS)
fun' = Lambda (Wise SOACS)
fun { lambdaBody :: Body (Wise SOACS)
lambdaBody = (Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun) { bodyResult :: [SubExp]
bodyResult = [SubExp]
ses' }
                      , lambdaReturnType :: [Type]
lambdaReturnType = [Type]
ts' }
       Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ PatternT (VarWisdom, Type)
Pattern (Lore (RuleM (Wise SOACS)))
pat' (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w (Lambda (Wise SOACS) -> ScremaForm (Wise SOACS)
forall lore. Bindable lore => Lambda lore -> ScremaForm lore
mapSOAC Lambda (Wise SOACS)
fun') [VName]
arrs
       [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))]
-> ((PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))
    -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))]
copies (((PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))
  -> RuleM (Wise SOACS) ())
 -> RuleM (Wise SOACS) ())
-> ((PatElemT (VarWisdom, Type), PatElemT (VarWisdom, Type))
    -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ \(PatElemT (VarWisdom, Type)
from,PatElemT (VarWisdom, Type)
to) ->
         if VName -> UsageTable -> Bool
UT.isConsumed (PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
to) UsageTable
used then
           Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ ([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)
to]) (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp (Wise SOACS)
forall lore. VName -> BasicOp lore
Copy (VName -> BasicOp (Wise SOACS)) -> VName -> BasicOp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
from
         else
           Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ ([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)
to]) (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Wise SOACS)
forall lore. SubExp -> BasicOp lore
SubExp (SubExp -> BasicOp (Wise SOACS)) -> SubExp -> BasicOp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
from
  where checkForDuplicates :: ([(a, b, a)], [(a, a)]) -> (a, b, a) -> ([(a, b, a)], [(a, a)])
checkForDuplicates ([(a, b, a)]
ses_ts_pes',[(a, a)]
copies) (a
se,b
t,a
pe)
          | Just (a
_,b
_,a
pe') <- ((a, b, a) -> Bool) -> [(a, b, a)] -> Maybe (a, b, a)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\(a
x,b
_,a
_) -> a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
se) [(a, b, a)]
ses_ts_pes' =
              -- This subexp has been returned before, producing the
              -- array pe'.
              ([(a, b, a)]
ses_ts_pes', (a
pe', a
pe) (a, a) -> [(a, a)] -> [(a, a)]
forall a. a -> [a] -> [a]
: [(a, a)]
copies)
          | Bool
otherwise = ([(a, b, a)]
ses_ts_pes' [(a, b, a)] -> [(a, b, a)] -> [(a, b, a)]
forall a. [a] -> [a] -> [a]
++ [(a
se,b
t,a
pe)], [(a, a)]
copies)
removeDuplicateMapOutput BottomUp (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

-- Mapping some operations becomes an extension of that operation.
mapOpToOp :: BottomUpRuleOp (Wise SOACS)

mapOpToOp :: RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
mapOpToOp (SymbolTable (Wise SOACS)
_, UsageTable
used) Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
aux1 Op (Wise SOACS)
e
  | Just (PatElemT (VarWisdom, Type)
map_pe, Certificates
cs, SubExp
w, BasicOp (Reshape ShapeChange SubExp
newshape VName
reshape_arr), [Param Type
p], [VName
arr]) <-
      PatternT (VarWisdom, Type)
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT (VarWisdom, Type), Certificates, SubExp,
      Exp (Wise SOACS), [Param Type], [VName])
forall attr.
PatternT attr
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
      [Param Type], [VName])
isMapWithOp PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat Op (Wise SOACS)
SOAC (Wise SOACS)
e,
    Param Type -> VName
forall attr. Param attr -> VName
paramName Param Type
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
reshape_arr,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName -> UsageTable -> Bool
UT.isConsumed (PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
map_pe) UsageTable
used = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      let redim :: DimChange SubExp
redim | Maybe [SubExp] -> Bool
forall a. Maybe a -> Bool
isJust (Maybe [SubExp] -> Bool) -> Maybe [SubExp] -> Bool
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> Maybe [SubExp]
forall d. ShapeChange d -> Maybe [d]
shapeCoercion ShapeChange SubExp
newshape = SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion SubExp
w
                | Bool
otherwise                       = SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
w
      Certificates -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux1 Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> Certificates
cs) (RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$
        BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp (Wise SOACS)
forall lore. ShapeChange SubExp -> VName -> BasicOp lore
Reshape (DimChange SubExp
redim DimChange SubExp -> ShapeChange SubExp -> ShapeChange SubExp
forall a. a -> [a] -> [a]
: ShapeChange SubExp
newshape) VName
arr

  | Just (PatElemT (VarWisdom, Type)
_, Certificates
cs, SubExp
_,
          BasicOp (Concat Int
d VName
arr [VName]
arrs SubExp
dw), [Param Type]
ps, VName
outer_arr : [VName]
outer_arrs) <-
      PatternT (VarWisdom, Type)
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT (VarWisdom, Type), Certificates, SubExp,
      Exp (Wise SOACS), [Param Type], [VName])
forall attr.
PatternT attr
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
      [Param Type], [VName])
isMapWithOp PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat Op (Wise SOACS)
SOAC (Wise SOACS)
e,
    (VName
arrVName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:[VName]
arrs) [VName] -> [VName] -> Bool
forall a. Eq a => a -> a -> Bool
== (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall attr. Param attr -> VName
paramName [Param Type]
ps =
      RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Certificates -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux1 Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> Certificates
cs) (RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$
      BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Int -> VName -> [VName] -> SubExp -> BasicOp (Wise SOACS)
forall lore. Int -> VName -> [VName] -> SubExp -> BasicOp lore
Concat (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) VName
outer_arr [VName]
outer_arrs SubExp
dw

  | Just (PatElemT (VarWisdom, Type)
map_pe, Certificates
cs, SubExp
_,
          BasicOp (Rearrange [Int]
perm VName
rearrange_arr), [Param Type
p], [VName
arr]) <-
      PatternT (VarWisdom, Type)
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT (VarWisdom, Type), Certificates, SubExp,
      Exp (Wise SOACS), [Param Type], [VName])
forall attr.
PatternT attr
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
      [Param Type], [VName])
isMapWithOp PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat Op (Wise SOACS)
SOAC (Wise SOACS)
e,
    Param Type -> VName
forall attr. Param attr -> VName
paramName Param Type
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
rearrange_arr,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName -> UsageTable -> Bool
UT.isConsumed (PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
map_pe) UsageTable
used =
      RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Certificates -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux1 Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> Certificates
cs) (RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$
      BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp (Wise SOACS)
forall lore. [Int] -> VName -> BasicOp lore
Rearrange (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+) [Int]
perm) VName
arr

  | Just (PatElemT (VarWisdom, Type)
map_pe, Certificates
cs, SubExp
_, BasicOp (Rotate [SubExp]
rots VName
rotate_arr), [Param Type
p], [VName
arr]) <-
      PatternT (VarWisdom, Type)
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT (VarWisdom, Type), Certificates, SubExp,
      Exp (Wise SOACS), [Param Type], [VName])
forall attr.
PatternT attr
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
      [Param Type], [VName])
isMapWithOp PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat Op (Wise SOACS)
SOAC (Wise SOACS)
e,
    Param Type -> VName
forall attr. Param attr -> VName
paramName Param Type
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
rotate_arr,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName -> UsageTable -> Bool
UT.isConsumed (PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
map_pe) UsageTable
used =
      RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Certificates -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux1 Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> Certificates
cs) (RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$
      BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> BasicOp (Wise SOACS)
forall lore. [SubExp] -> VName -> BasicOp lore
Rotate (IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
0 SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
rots) VName
arr

mapOpToOp BottomUp (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

isMapWithOp :: PatternT attr
            -> SOAC (Wise SOACS)
            -> Maybe (PatElemT attr, Certificates, SubExp,
                      AST.Exp (Wise SOACS), [Param Type], [VName])
isMapWithOp :: PatternT attr
-> SOAC (Wise SOACS)
-> Maybe
     (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
      [Param Type], [VName])
isMapWithOp PatternT attr
pat SOAC (Wise SOACS)
e
  | Pattern [] [PatElemT attr
map_pe] <- PatternT attr
pat,
    Screma SubExp
w ScremaForm (Wise SOACS)
form [VName]
arrs <- SOAC (Wise SOACS)
e,
    Just Lambda (Wise SOACS)
map_lam <- ScremaForm (Wise SOACS) -> Maybe (Lambda (Wise SOACS))
forall lore. ScremaForm lore -> Maybe (Lambda lore)
isMapSOAC ScremaForm (Wise SOACS)
form,
    [Let (Pattern [] [PatElemT (LetAttr (Wise SOACS))
pe]) StmAux (ExpAttr (Wise SOACS))
aux2 Exp (Wise SOACS)
e'] <-
      Stms (Wise SOACS) -> [Stm (Wise SOACS)]
forall lore. Stms lore -> [Stm lore]
stmsToList (Stms (Wise SOACS) -> [Stm (Wise SOACS)])
-> Stms (Wise SOACS) -> [Stm (Wise SOACS)]
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> Stms (Wise SOACS)
forall lore. BodyT lore -> Stms lore
bodyStms (Body (Wise SOACS) -> Stms (Wise SOACS))
-> Body (Wise SOACS) -> Stms (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam,
    [Var VName
r] <- Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam,
    VName
r VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
PatElemT (LetAttr (Wise SOACS))
pe =
      (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
 [Param Type], [VName])
-> Maybe
     (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
      [Param Type], [VName])
forall a. a -> Maybe a
Just (PatElemT attr
map_pe, StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux2, SubExp
w, Exp (Wise SOACS)
e', Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
map_lam, [VName]
arrs)
  | Bool
otherwise = Maybe
  (PatElemT attr, Certificates, SubExp, Exp (Wise SOACS),
   [Param Type], [VName])
forall a. Maybe a
Nothing

-- | Some of the results of a reduction (or really: Redomap) may be
-- dead.  We remove them here.  The trick is that we need to look at
-- the data dependencies to see that the "dead" result is not
-- actually used for computing one of the live ones.
removeDeadReduction :: BottomUpRuleOp (Wise SOACS)
removeDeadReduction :: RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDeadReduction (SymbolTable (Wise SOACS)
_, UsageTable
used) Pattern (Wise SOACS)
pat (StmAux Certificates
cs ExpAttr (Wise SOACS)
_) (Screma w form arrs)
  | Just ([Reduce Commutativity
comm Lambda (Wise SOACS)
redlam [SubExp]
nes], Lambda (Wise SOACS)
maplam) <- ScremaForm (Wise SOACS)
-> Maybe ([Reduce (Wise SOACS)], Lambda (Wise SOACS))
forall lore. ScremaForm lore -> Maybe ([Reduce lore], Lambda lore)
isRedomapSOAC ScremaForm (Wise SOACS)
form,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> UsageTable -> Bool
`UT.used` UsageTable
used) ([VName] -> Bool) -> [VName] -> Bool
forall a b. (a -> b) -> a -> b
$ PatternT (VarWisdom, Type) -> [VName]
forall attr. PatternT attr -> [VName]
patternNames PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat, -- Quick/cheap check

    let ([PatElemT (VarWisdom, Type)]
red_pes, [PatElemT (VarWisdom, Type)]
map_pes) = Int
-> [PatElemT (VarWisdom, Type)]
-> ([PatElemT (VarWisdom, Type)], [PatElemT (VarWisdom, Type)])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([PatElemT (VarWisdom, Type)]
 -> ([PatElemT (VarWisdom, Type)], [PatElemT (VarWisdom, Type)]))
-> [PatElemT (VarWisdom, Type)]
-> ([PatElemT (VarWisdom, Type)], [PatElemT (VarWisdom, Type)])
forall a b. (a -> b) -> a -> b
$ PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternElements PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat,
    let redlam_deps :: Dependencies
redlam_deps = Body (Wise SOACS) -> Dependencies
forall lore. Attributes lore => Body lore -> Dependencies
dataDependencies (Body (Wise SOACS) -> Dependencies)
-> Body (Wise SOACS) -> Dependencies
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
redlam,
    let redlam_res :: [SubExp]
redlam_res = Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
redlam,
    let redlam_params :: [LParam (Wise SOACS)]
redlam_params = Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
redlam,
    let used_after :: [Param Type]
used_after = ((PatElemT (VarWisdom, Type), Param Type) -> Param Type)
-> [(PatElemT (VarWisdom, Type), Param Type)] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (PatElemT (VarWisdom, Type), Param Type) -> Param Type
forall a b. (a, b) -> b
snd ([(PatElemT (VarWisdom, Type), Param Type)] -> [Param Type])
-> [(PatElemT (VarWisdom, Type), Param Type)] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ ((PatElemT (VarWisdom, Type), Param Type) -> Bool)
-> [(PatElemT (VarWisdom, Type), Param Type)]
-> [(PatElemT (VarWisdom, Type), Param Type)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> UsageTable -> Bool
`UT.used` UsageTable
used) (VName -> Bool)
-> ((PatElemT (VarWisdom, Type), Param Type) -> VName)
-> (PatElemT (VarWisdom, Type), Param Type)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName (PatElemT (VarWisdom, Type) -> VName)
-> ((PatElemT (VarWisdom, Type), Param Type)
    -> PatElemT (VarWisdom, Type))
-> (PatElemT (VarWisdom, Type), Param Type)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PatElemT (VarWisdom, Type), Param Type)
-> PatElemT (VarWisdom, Type)
forall a b. (a, b) -> a
fst) ([(PatElemT (VarWisdom, Type), Param Type)]
 -> [(PatElemT (VarWisdom, Type), Param Type)])
-> [(PatElemT (VarWisdom, Type), Param Type)]
-> [(PatElemT (VarWisdom, Type), Param Type)]
forall a b. (a -> b) -> a -> b
$
                     [PatElemT (VarWisdom, Type)]
-> [Param Type] -> [(PatElemT (VarWisdom, Type), Param Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatElemT (VarWisdom, Type)]
red_pes [Param Type]
[LParam (Wise SOACS)]
redlam_params,
    let necessary :: Names
necessary = (Param Type -> Bool)
-> [(Param Type, SubExp)] -> Dependencies -> Names
forall attr.
(Param attr -> Bool)
-> [(Param attr, SubExp)] -> Dependencies -> Names
findNecessaryForReturned (Param Type -> [Param Type] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Param Type]
used_after)
                    ([Param Type] -> [SubExp] -> [(Param Type, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
[LParam (Wise SOACS)]
redlam_params ([SubExp] -> [(Param Type, SubExp)])
-> [SubExp] -> [(Param Type, SubExp)]
forall a b. (a -> b) -> a -> b
$ [SubExp]
redlam_res [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> [SubExp]
redlam_res) Dependencies
redlam_deps,
    let alive_mask :: [Bool]
alive_mask = (Param Type -> Bool) -> [Param Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> Names -> Bool
`nameIn` Names
necessary) (VName -> Bool) -> (Param Type -> VName) -> Param Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall attr. Param attr -> VName
paramName) [Param Type]
[LParam (Wise SOACS)]
redlam_params,

    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool) -> [Bool] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
==Bool
True) [Bool]
alive_mask = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do

  let fixDeadToNeutral :: Bool -> a -> Maybe a
fixDeadToNeutral Bool
lives a
ne = if Bool
lives then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just a
ne
      dead_fix :: [Maybe SubExp]
dead_fix = (Bool -> SubExp -> Maybe SubExp)
-> [Bool] -> [SubExp] -> [Maybe SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Bool -> SubExp -> Maybe SubExp
forall a. Bool -> a -> Maybe a
fixDeadToNeutral [Bool]
alive_mask [SubExp]
nes
      ([PatElemT (VarWisdom, Type)]
used_red_pes, [Param Type]
_, [SubExp]
used_nes) =
        [(PatElemT (VarWisdom, Type), Param Type, SubExp)]
-> ([PatElemT (VarWisdom, Type)], [Param Type], [SubExp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(PatElemT (VarWisdom, Type), Param Type, SubExp)]
 -> ([PatElemT (VarWisdom, Type)], [Param Type], [SubExp]))
-> [(PatElemT (VarWisdom, Type), Param Type, SubExp)]
-> ([PatElemT (VarWisdom, Type)], [Param Type], [SubExp])
forall a b. (a -> b) -> a -> b
$ ((PatElemT (VarWisdom, Type), Param Type, SubExp) -> Bool)
-> [(PatElemT (VarWisdom, Type), Param Type, SubExp)]
-> [(PatElemT (VarWisdom, Type), Param Type, SubExp)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(PatElemT (VarWisdom, Type)
_,Param Type
x,SubExp
_) -> Param Type -> VName
forall attr. Param attr -> VName
paramName Param Type
x VName -> Names -> Bool
`nameIn` Names
necessary) ([(PatElemT (VarWisdom, Type), Param Type, SubExp)]
 -> [(PatElemT (VarWisdom, Type), Param Type, SubExp)])
-> [(PatElemT (VarWisdom, Type), Param Type, SubExp)]
-> [(PatElemT (VarWisdom, Type), Param Type, SubExp)]
forall a b. (a -> b) -> a -> b
$
        [PatElemT (VarWisdom, Type)]
-> [Param Type]
-> [SubExp]
-> [(PatElemT (VarWisdom, Type), Param Type, SubExp)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [PatElemT (VarWisdom, Type)]
red_pes [Param Type]
[LParam (Wise SOACS)]
redlam_params [SubExp]
nes

  let maplam' :: Lambda (Wise SOACS)
maplam' = [Bool] -> Lambda (Wise SOACS) -> Lambda (Wise SOACS)
forall lore. [Bool] -> Lambda lore -> Lambda lore
removeLambdaResults (Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) [Bool]
alive_mask) Lambda (Wise SOACS)
maplam
  Lambda (Wise SOACS)
redlam' <- [Bool] -> Lambda (Wise SOACS) -> Lambda (Wise SOACS)
forall lore. [Bool] -> Lambda lore -> Lambda lore
removeLambdaResults (Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) [Bool]
alive_mask) (Lambda (Wise SOACS) -> Lambda (Wise SOACS))
-> RuleM (Wise SOACS) (Lambda (Wise SOACS))
-> RuleM (Wise SOACS) (Lambda (Wise SOACS))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda (Lore (RuleM (Wise SOACS)))
-> [Maybe SubExp]
-> RuleM (Wise SOACS) (Lambda (Lore (RuleM (Wise SOACS))))
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m), BinderOps (Lore m)) =>
Lambda (Lore m) -> [Maybe SubExp] -> m (Lambda (Lore m))
fixLambdaParams Lambda (Lore (RuleM (Wise SOACS)))
Lambda (Wise SOACS)
redlam ([Maybe SubExp]
dead_fix[Maybe SubExp] -> [Maybe SubExp] -> [Maybe SubExp]
forall a. [a] -> [a] -> [a]
++[Maybe SubExp]
dead_fix)

  Certificates -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ ([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] ([PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type))
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall a b. (a -> b) -> a -> b
$ [PatElemT (VarWisdom, Type)]
used_red_pes [PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> [PatElemT (VarWisdom, Type)]
forall a. [a] -> [a] -> [a]
++ [PatElemT (VarWisdom, Type)]
map_pes) (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$
    Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w ([Reduce (Wise SOACS)]
-> Lambda (Wise SOACS) -> ScremaForm (Wise SOACS)
forall lore.
Bindable lore =>
[Reduce lore] -> Lambda lore -> ScremaForm lore
redomapSOAC [Commutativity
-> Lambda (Wise SOACS) -> [SubExp] -> Reduce (Wise SOACS)
forall lore.
Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
Reduce Commutativity
comm Lambda (Wise SOACS)
redlam' [SubExp]
used_nes] Lambda (Wise SOACS)
maplam') [VName]
arrs

removeDeadReduction BottomUp (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

-- | If we are writing to an array that is never used, get rid of it.
removeDeadWrite :: BottomUpRuleOp (Wise SOACS)
removeDeadWrite :: RuleOp (Wise SOACS) (BottomUp (Wise SOACS))
removeDeadWrite (SymbolTable (Wise SOACS)
_, UsageTable
used) Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Scatter w fun arrs dests) =
  let ([SubExp]
i_ses, [SubExp]
v_ses) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([(SubExp, Int, VName)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(SubExp, Int, VName)]
dests) ([SubExp] -> ([SubExp], [SubExp]))
-> [SubExp] -> ([SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp]) -> Body (Wise SOACS) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun
      ([Type]
i_ts, [Type]
v_ts) = Int -> [Type] -> ([Type], [Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([(SubExp, Int, VName)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(SubExp, Int, VName)]
dests) ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda (Wise SOACS)
fun
      isUsed :: (PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
 (SubExp, Int, VName))
-> Bool
isUsed (PatElemT (VarWisdom, Type)
bindee, SubExp
_, SubExp
_, Type
_, Type
_, (SubExp, Int, VName)
_) = (VName -> UsageTable -> Bool
`UT.used` UsageTable
used) (VName -> Bool) -> VName -> Bool
forall a b. (a -> b) -> a -> b
$ PatElemT (VarWisdom, Type) -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT (VarWisdom, Type)
bindee
      ([PatElemT (VarWisdom, Type)]
pat', [SubExp]
i_ses', [SubExp]
v_ses', [Type]
i_ts', [Type]
v_ts', [(SubExp, Int, VName)]
dests') =
        [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
  (SubExp, Int, VName))]
-> ([PatElemT (VarWisdom, Type)], [SubExp], [SubExp], [Type],
    [Type], [(SubExp, Int, VName)])
forall a b c d e f.
[(a, b, c, d, e, f)] -> ([a], [b], [c], [d], [e], [f])
unzip6 ([(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
   (SubExp, Int, VName))]
 -> ([PatElemT (VarWisdom, Type)], [SubExp], [SubExp], [Type],
     [Type], [(SubExp, Int, VName)]))
-> [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
     (SubExp, Int, VName))]
-> ([PatElemT (VarWisdom, Type)], [SubExp], [SubExp], [Type],
    [Type], [(SubExp, Int, VName)])
forall a b. (a -> b) -> a -> b
$ ((PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
  (SubExp, Int, VName))
 -> Bool)
-> [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
     (SubExp, Int, VName))]
-> [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
     (SubExp, Int, VName))]
forall a. (a -> Bool) -> [a] -> [a]
filter (PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
 (SubExp, Int, VName))
-> Bool
isUsed ([(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
   (SubExp, Int, VName))]
 -> [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
      (SubExp, Int, VName))])
-> [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
     (SubExp, Int, VName))]
-> [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
     (SubExp, Int, VName))]
forall a b. (a -> b) -> a -> b
$
        [PatElemT (VarWisdom, Type)]
-> [SubExp]
-> [SubExp]
-> [Type]
-> [Type]
-> [(SubExp, Int, VName)]
-> [(PatElemT (VarWisdom, Type), SubExp, SubExp, Type, Type,
     (SubExp, Int, VName))]
forall a b c d e f.
[a] -> [b] -> [c] -> [d] -> [e] -> [f] -> [(a, b, c, d, e, f)]
zip6 (PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternElements PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat) [SubExp]
i_ses [SubExp]
v_ses [Type]
i_ts [Type]
v_ts [(SubExp, Int, VName)]
dests
      fun' :: Lambda (Wise SOACS)
fun' = Lambda (Wise SOACS)
fun { lambdaBody :: Body (Wise SOACS)
lambdaBody = (Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
fun) { bodyResult :: [SubExp]
bodyResult = [SubExp]
i_ses' [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
v_ses' }
                 , lambdaReturnType :: [Type]
lambdaReturnType = [Type]
i_ts' [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
v_ts'
                 }
  in if PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat PatternT (VarWisdom, Type) -> PatternT (VarWisdom, Type) -> Bool
forall a. Eq a => a -> a -> Bool
/= [PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)]
pat'
     then RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ ([PatElemT (VarWisdom, Type)]
-> [PatElemT (VarWisdom, Type)] -> PatternT (VarWisdom, Type)
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [PatElemT (VarWisdom, Type)]
pat') (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp
-> Lambda (Wise SOACS)
-> [VName]
-> [(SubExp, Int, VName)]
-> SOAC (Wise SOACS)
forall lore.
SubExp
-> Lambda lore -> [VName] -> [(SubExp, Int, VName)] -> SOAC lore
Scatter SubExp
w Lambda (Wise SOACS)
fun' [VName]
arrs [(SubExp, Int, VName)]
dests'
     else Rule (Wise SOACS)
forall lore. Rule lore
Skip
removeDeadWrite BottomUp (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

-- handles now concatenation of more than two arrays
fuseConcatScatter :: TopDownRuleOp (Wise SOACS)
fuseConcatScatter :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
fuseConcatScatter SymbolTable (Wise SOACS)
vtable Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Scatter _ fun arrs dests)
  | Just (ws :: [SubExp]
ws@(SubExp
w':[SubExp]
_), [[VName]]
xss, [Certificates]
css) <- [(SubExp, [VName], Certificates)]
-> ([SubExp], [[VName]], [Certificates])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(SubExp, [VName], Certificates)]
 -> ([SubExp], [[VName]], [Certificates]))
-> Maybe [(SubExp, [VName], Certificates)]
-> Maybe ([SubExp], [[VName]], [Certificates])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> Maybe (SubExp, [VName], Certificates))
-> [VName] -> Maybe [(SubExp, [VName], Certificates)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> Maybe (SubExp, [VName], Certificates)
isConcat [VName]
arrs,
    [[VName]]
xivs <- [[VName]] -> [[VName]]
forall a. [[a]] -> [[a]]
transpose [[VName]]
xss,
    (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (SubExp
w'SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
==) [SubExp]
ws = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      let r :: Int
r = [[VName]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[VName]]
xivs
      [Lambda (Wise SOACS)]
fun2s <- (Int -> RuleM (Wise SOACS) (Lambda (Wise SOACS)))
-> [Int] -> RuleM (Wise SOACS) [Lambda (Wise SOACS)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Int
_ -> Lambda (Wise SOACS) -> RuleM (Wise SOACS) (Lambda (Wise SOACS))
forall lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Lambda lore -> m (Lambda lore)
renameLambda Lambda (Wise SOACS)
fun) [Int
1 .. Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
      let fun_n :: Int
fun_n = [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Type] -> Int) -> [Type] -> Int
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda (Wise SOACS)
fun
          ([[SubExp]]
fun_is, [[SubExp]]
fun_vs) = [([SubExp], [SubExp])] -> ([[SubExp]], [[SubExp]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([SubExp], [SubExp])] -> ([[SubExp]], [[SubExp]]))
-> [([SubExp], [SubExp])] -> ([[SubExp]], [[SubExp]])
forall a b. (a -> b) -> a -> b
$ (Lambda (Wise SOACS) -> ([SubExp], [SubExp]))
-> [Lambda (Wise SOACS)] -> [([SubExp], [SubExp])]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
fun_n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) ([SubExp] -> ([SubExp], [SubExp]))
-> (Lambda (Wise SOACS) -> [SubExp])
-> Lambda (Wise SOACS)
-> ([SubExp], [SubExp])
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                             Body (Wise SOACS) -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (Body (Wise SOACS) -> [SubExp])
-> (Lambda (Wise SOACS) -> Body (Wise SOACS))
-> Lambda (Wise SOACS)
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody ) (Lambda (Wise SOACS)
funLambda (Wise SOACS)
-> [Lambda (Wise SOACS)] -> [Lambda (Wise SOACS)]
forall a. a -> [a] -> [a]
:[Lambda (Wise SOACS)]
fun2s)
          ([[Type]]
its, [[Type]]
vts) = [([Type], [Type])] -> ([[Type]], [[Type]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([Type], [Type])] -> ([[Type]], [[Type]]))
-> [([Type], [Type])] -> ([[Type]], [[Type]])
forall a b. (a -> b) -> a -> b
$ Int -> ([Type], [Type]) -> [([Type], [Type])]
forall a. Int -> a -> [a]
replicate Int
r (([Type], [Type]) -> [([Type], [Type])])
-> ([Type], [Type]) -> [([Type], [Type])]
forall a b. (a -> b) -> a -> b
$
                       Int -> [Type] -> ([Type], [Type])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
fun_n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda (Wise SOACS)
fun
          new_stmts :: Stms (Wise SOACS)
new_stmts  = [Stms (Wise SOACS)] -> Stms (Wise SOACS)
forall a. Monoid a => [a] -> a
mconcat ([Stms (Wise SOACS)] -> Stms (Wise SOACS))
-> [Stms (Wise SOACS)] -> Stms (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ (Lambda (Wise SOACS) -> Stms (Wise SOACS))
-> [Lambda (Wise SOACS)] -> [Stms (Wise SOACS)]
forall a b. (a -> b) -> [a] -> [b]
map (Body (Wise SOACS) -> Stms (Wise SOACS)
forall lore. BodyT lore -> Stms lore
bodyStms (Body (Wise SOACS) -> Stms (Wise SOACS))
-> (Lambda (Wise SOACS) -> Body (Wise SOACS))
-> Lambda (Wise SOACS)
-> Stms (Wise SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody) (Lambda (Wise SOACS)
funLambda (Wise SOACS)
-> [Lambda (Wise SOACS)] -> [Lambda (Wise SOACS)]
forall a. a -> [a] -> [a]
:[Lambda (Wise SOACS)]
fun2s)
      let fun' :: Lambda (Wise SOACS)
fun' = Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
Lambda
                 { lambdaParams :: [LParam (Wise SOACS)]
lambdaParams = [[Param Type]] -> [Param Type]
forall a. Monoid a => [a] -> a
mconcat ([[Param Type]] -> [Param Type]) -> [[Param Type]] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ (Lambda (Wise SOACS) -> [Param Type])
-> [Lambda (Wise SOACS)] -> [[Param Type]]
forall a b. (a -> b) -> [a] -> [b]
map Lambda (Wise SOACS) -> [Param Type]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams (Lambda (Wise SOACS)
funLambda (Wise SOACS)
-> [Lambda (Wise SOACS)] -> [Lambda (Wise SOACS)]
forall a. a -> [a] -> [a]
:[Lambda (Wise SOACS)]
fun2s)
                 , lambdaBody :: Body (Wise SOACS)
lambdaBody = Stms (Wise SOACS) -> [SubExp] -> Body (Wise SOACS)
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms (Wise SOACS)
new_stmts ([SubExp] -> Body (Wise SOACS)) -> [SubExp] -> Body (Wise SOACS)
forall a b. (a -> b) -> a -> b
$
                                [[SubExp]] -> [SubExp]
forall a. [[a]] -> [a]
mix [[SubExp]]
fun_is [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> [[SubExp]] -> [SubExp]
forall a. [[a]] -> [a]
mix [[SubExp]]
fun_vs
                 , lambdaReturnType :: [Type]
lambdaReturnType = [[Type]] -> [Type]
forall a. [[a]] -> [a]
mix [[Type]]
its [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> [[Type]] -> [Type]
forall a. [[a]] -> [a]
mix [[Type]]
vts
                 }
      Certificates -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying ([Certificates] -> Certificates
forall a. Monoid a => [a] -> a
mconcat [Certificates]
css) (RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$
        Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp
-> Lambda (Wise SOACS)
-> [VName]
-> [(SubExp, Int, VName)]
-> SOAC (Wise SOACS)
forall lore.
SubExp
-> Lambda lore -> [VName] -> [(SubExp, Int, VName)] -> SOAC lore
Scatter SubExp
w' Lambda (Wise SOACS)
fun' ([[VName]] -> [VName]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[VName]]
xivs) ([(SubExp, Int, VName)] -> SOAC (Wise SOACS))
-> [(SubExp, Int, VName)] -> SOAC (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ ((SubExp, Int, VName) -> (SubExp, Int, VName))
-> [(SubExp, Int, VName)] -> [(SubExp, Int, VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> (SubExp, Int, VName) -> (SubExp, Int, VName)
forall b a c. Num b => b -> (a, b, c) -> (a, b, c)
incWrites Int
r) [(SubExp, Int, VName)]
dests
  where sizeOf :: VName -> Maybe SubExp
        sizeOf :: VName -> Maybe SubExp
sizeOf VName
x = Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 (Type -> SubExp)
-> (Entry (Wise SOACS) -> Type) -> Entry (Wise SOACS) -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry (Wise SOACS) -> Type
forall lore. Attributes lore => Entry lore -> Type
ST.entryType (Entry (Wise SOACS) -> SubExp)
-> Maybe (Entry (Wise SOACS)) -> Maybe SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SymbolTable (Wise SOACS) -> Maybe (Entry (Wise SOACS))
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
ST.lookup VName
x SymbolTable (Wise SOACS)
vtable
        mix :: [[a]] -> [a]
mix = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> ([[a]] -> [[a]]) -> [[a]] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose
        incWrites :: b -> (a, b, c) -> (a, b, c)
incWrites b
r (a
w, b
n, c
a) = (a
w, b
nb -> b -> b
forall a. Num a => a -> a -> a
*b
r, c
a) -- ToDO: is it (n*r) or (n+r-1)??
        isConcat :: VName -> Maybe (SubExp, [VName], Certificates)
isConcat VName
v = case VName
-> SymbolTable (Wise SOACS)
-> Maybe (Exp (Wise SOACS), Certificates)
forall lore.
VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
ST.lookupExp VName
v SymbolTable (Wise SOACS)
vtable of
          Just (BasicOp (Concat Int
0 VName
x [VName]
ys SubExp
_), Certificates
cs) -> do
            SubExp
x_w <- VName -> Maybe SubExp
sizeOf VName
x
            [SubExp]
y_ws<- (VName -> Maybe SubExp) -> [VName] -> Maybe [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> Maybe SubExp
sizeOf [VName]
ys
            Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ (SubExp -> Bool) -> [SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (SubExp
x_wSubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
==) [SubExp]
y_ws
            (SubExp, [VName], Certificates)
-> Maybe (SubExp, [VName], Certificates)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
x_w, VName
xVName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:[VName]
ys, Certificates
cs)
          Just (BasicOp (Reshape ShapeChange SubExp
reshape VName
arr), Certificates
cs) -> do
            Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Maybe [SubExp] -> Bool
forall a. Maybe a -> Bool
isJust (Maybe [SubExp] -> Bool) -> Maybe [SubExp] -> Bool
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> Maybe [SubExp]
forall d. ShapeChange d -> Maybe [d]
shapeCoercion ShapeChange SubExp
reshape
            (SubExp
a, [VName]
b, Certificates
cs') <- VName -> Maybe (SubExp, [VName], Certificates)
isConcat VName
arr
            (SubExp, [VName], Certificates)
-> Maybe (SubExp, [VName], Certificates)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
a, [VName]
b, Certificates
cs Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> Certificates
cs')
          Maybe (Exp (Wise SOACS), Certificates)
_ -> Maybe (SubExp, [VName], Certificates)
forall a. Maybe a
Nothing

fuseConcatScatter SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

simplifyClosedFormReduce :: TopDownRuleOp (Wise SOACS)
simplifyClosedFormReduce :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
simplifyClosedFormReduce SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma (Constant w) form _)
  | Just [SubExp]
nes <- (Reduce (Wise SOACS) -> [SubExp])
-> [Reduce (Wise SOACS)] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Reduce (Wise SOACS) -> [SubExp]
forall lore. Reduce lore -> [SubExp]
redNeutral ([Reduce (Wise SOACS)] -> [SubExp])
-> (([Reduce (Wise SOACS)], Lambda (Wise SOACS))
    -> [Reduce (Wise SOACS)])
-> ([Reduce (Wise SOACS)], Lambda (Wise SOACS))
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Reduce (Wise SOACS)], Lambda (Wise SOACS))
-> [Reduce (Wise SOACS)]
forall a b. (a, b) -> a
fst (([Reduce (Wise SOACS)], Lambda (Wise SOACS)) -> [SubExp])
-> Maybe ([Reduce (Wise SOACS)], Lambda (Wise SOACS))
-> Maybe [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScremaForm (Wise SOACS)
-> Maybe ([Reduce (Wise SOACS)], Lambda (Wise SOACS))
forall lore. ScremaForm lore -> Maybe ([Reduce lore], Lambda lore)
isRedomapSOAC ScremaForm (Wise SOACS)
form,
    PrimValue -> Bool
zeroIsh PrimValue
w =
      RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [(VName, SubExp)]
-> ((VName, SubExp) -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (PatternT (VarWisdom, Type) -> [VName]
forall attr. PatternT attr -> [VName]
patternNames PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat) [SubExp]
nes) (((VName, SubExp) -> RuleM (Wise SOACS) ())
 -> RuleM (Wise SOACS) ())
-> ((VName, SubExp) -> RuleM (Wise SOACS) ())
-> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
ne) ->
      [VName] -> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [VName
v] (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Wise SOACS)
forall lore. SubExp -> BasicOp lore
SubExp SubExp
ne
simplifyClosedFormReduce SymbolTable (Wise SOACS)
vtable Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma _ form arrs)
  | Just [Reduce Commutativity
_ Lambda (Wise SOACS)
red_fun [SubExp]
nes] <- ScremaForm (Wise SOACS) -> Maybe [Reduce (Wise SOACS)]
forall lore. ScremaForm lore -> Maybe [Reduce lore]
isReduceSOAC ScremaForm (Wise SOACS)
form =
      RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VarLookup (Wise SOACS)
-> Pattern (Wise SOACS)
-> Lambda (Wise SOACS)
-> [SubExp]
-> [VName]
-> RuleM (Wise SOACS) ()
forall lore.
(Attributes lore, BinderOps lore) =>
VarLookup lore
-> Pattern lore
-> Lambda lore
-> [SubExp]
-> [VName]
-> RuleM lore ()
foldClosedForm (VName
-> SymbolTable (Wise SOACS)
-> Maybe (Exp (Wise SOACS), Certificates)
forall lore.
VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
`ST.lookupExp` SymbolTable (Wise SOACS)
vtable) Pattern (Wise SOACS)
pat Lambda (Wise SOACS)
red_fun [SubExp]
nes [VName]
arrs
simplifyClosedFormReduce SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

-- For now we just remove singleton SOACs.
simplifyKnownIterationSOAC :: TopDownRuleOp (Wise SOACS)
simplifyKnownIterationSOAC :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
simplifyKnownIterationSOAC SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma (Constant k)
                                    (ScremaForm (scan_lam, scan_nes) reds map_lam)
                                    arrs)
  | PrimValue -> Bool
oneIsh PrimValue
k = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      (Param Type -> VName -> RuleM (Wise SOACS) ())
-> [Param Type] -> [VName] -> RuleM (Wise SOACS) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Param Type -> VName -> RuleM (Wise SOACS) ()
forall (m :: * -> *) attr.
MonadBinder m =>
Param attr -> VName -> m ()
bindMapParam (Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
map_lam) [VName]
arrs
      ([SubExp]
to_scan, [SubExp]
to_red, [SubExp]
map_res) <- Int -> Int -> [SubExp] -> ([SubExp], [SubExp], [SubExp])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
scan_nes) ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
red_nes) ([SubExp] -> ([SubExp], [SubExp], [SubExp]))
-> RuleM (Wise SOACS) [SubExp]
-> RuleM (Wise SOACS) ([SubExp], [SubExp], [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                                    Body (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind (Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam)
      [SubExp]
scan_res <- Lambda (Lore (RuleM (Wise SOACS)))
-> [RuleM (Wise SOACS) (Exp (Lore (RuleM (Wise SOACS))))]
-> RuleM (Wise SOACS) [SubExp]
forall (m :: * -> *).
MonadBinder m =>
Lambda (Lore m) -> [m (Exp (Lore m))] -> m [SubExp]
eLambda Lambda (Lore (RuleM (Wise SOACS)))
Lambda (Wise SOACS)
scan_lam ([RuleM (Wise SOACS) (Exp (Lore (RuleM (Wise SOACS))))]
 -> RuleM (Wise SOACS) [SubExp])
-> [RuleM (Wise SOACS) (Exp (Lore (RuleM (Wise SOACS))))]
-> RuleM (Wise SOACS) [SubExp]
forall a b. (a -> b) -> a -> b
$ (SubExp -> RuleM (Wise SOACS) (Exp (Wise SOACS)))
-> [SubExp] -> [RuleM (Wise SOACS) (Exp (Wise SOACS))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> RuleM (Wise SOACS) (Exp (Wise SOACS))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp ([SubExp] -> [RuleM (Wise SOACS) (Exp (Wise SOACS))])
-> [SubExp] -> [RuleM (Wise SOACS) (Exp (Wise SOACS))]
forall a b. (a -> b) -> a -> b
$ [SubExp]
scan_nes [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
to_scan
      [SubExp]
red_res <- Lambda (Lore (RuleM (Wise SOACS)))
-> [RuleM (Wise SOACS) (Exp (Lore (RuleM (Wise SOACS))))]
-> RuleM (Wise SOACS) [SubExp]
forall (m :: * -> *).
MonadBinder m =>
Lambda (Lore m) -> [m (Exp (Lore m))] -> m [SubExp]
eLambda Lambda (Lore (RuleM (Wise SOACS)))
Lambda (Wise SOACS)
red_lam ([RuleM (Wise SOACS) (Exp (Lore (RuleM (Wise SOACS))))]
 -> RuleM (Wise SOACS) [SubExp])
-> [RuleM (Wise SOACS) (Exp (Lore (RuleM (Wise SOACS))))]
-> RuleM (Wise SOACS) [SubExp]
forall a b. (a -> b) -> a -> b
$ (SubExp -> RuleM (Wise SOACS) (Exp (Wise SOACS)))
-> [SubExp] -> [RuleM (Wise SOACS) (Exp (Wise SOACS))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> RuleM (Wise SOACS) (Exp (Wise SOACS))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp ([SubExp] -> [RuleM (Wise SOACS) (Exp (Wise SOACS))])
-> [SubExp] -> [RuleM (Wise SOACS) (Exp (Wise SOACS))]
forall a b. (a -> b) -> a -> b
$ [SubExp]
red_nes [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
to_red

      (PatElemT (VarWisdom, Type) -> SubExp -> RuleM (Wise SOACS) ())
-> [PatElemT (VarWisdom, Type)]
-> [SubExp]
-> RuleM (Wise SOACS) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ PatElemT (VarWisdom, Type) -> SubExp -> RuleM (Wise SOACS) ()
forall (m :: * -> *) attr.
(MonadBinder m, Typed attr) =>
PatElemT attr -> SubExp -> m ()
bindArrayResult [PatElemT (VarWisdom, Type)]
scan_pes [SubExp]
scan_res
      (PatElemT (VarWisdom, Type) -> SubExp -> RuleM (Wise SOACS) ())
-> [PatElemT (VarWisdom, Type)]
-> [SubExp]
-> RuleM (Wise SOACS) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ PatElemT (VarWisdom, Type) -> SubExp -> RuleM (Wise SOACS) ()
forall (m :: * -> *) attr.
MonadBinder m =>
PatElemT attr -> SubExp -> m ()
bindResult [PatElemT (VarWisdom, Type)]
red_pes [SubExp]
red_res
      (PatElemT (VarWisdom, Type) -> SubExp -> RuleM (Wise SOACS) ())
-> [PatElemT (VarWisdom, Type)]
-> [SubExp]
-> RuleM (Wise SOACS) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ PatElemT (VarWisdom, Type) -> SubExp -> RuleM (Wise SOACS) ()
forall (m :: * -> *) attr.
(MonadBinder m, Typed attr) =>
PatElemT attr -> SubExp -> m ()
bindArrayResult [PatElemT (VarWisdom, Type)]
map_pes [SubExp]
map_res

        where (Reduce Commutativity
_ Lambda (Wise SOACS)
red_lam [SubExp]
red_nes) = [Reduce (Wise SOACS)] -> Reduce (Wise SOACS)
forall lore. Bindable lore => [Reduce lore] -> Reduce lore
singleReduce [Reduce (Wise SOACS)]
reds
              ([PatElemT (VarWisdom, Type)]
scan_pes, [PatElemT (VarWisdom, Type)]
red_pes, [PatElemT (VarWisdom, Type)]
map_pes) = Int
-> Int
-> [PatElemT (VarWisdom, Type)]
-> ([PatElemT (VarWisdom, Type)], [PatElemT (VarWisdom, Type)],
    [PatElemT (VarWisdom, Type)])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
scan_nes) ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
red_nes) ([PatElemT (VarWisdom, Type)]
 -> ([PatElemT (VarWisdom, Type)], [PatElemT (VarWisdom, Type)],
     [PatElemT (VarWisdom, Type)]))
-> [PatElemT (VarWisdom, Type)]
-> ([PatElemT (VarWisdom, Type)], [PatElemT (VarWisdom, Type)],
    [PatElemT (VarWisdom, Type)])
forall a b. (a -> b) -> a -> b
$
                                             PatternT (VarWisdom, Type) -> [PatElemT (VarWisdom, Type)]
forall attr. PatternT attr -> [PatElemT attr]
patternElements PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat
              bindMapParam :: Param attr -> VName -> m ()
bindMapParam Param attr
p VName
a = do
                Type
a_t <- VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
a
                [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [Param attr -> VName
forall attr. Param attr -> VName
paramName Param attr
p] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$
                  BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> BasicOp (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp (Lore m)
forall lore. VName -> Slice SubExp -> BasicOp lore
Index VName
a (Slice SubExp -> BasicOp (Lore m))
-> Slice SubExp -> BasicOp (Lore m)
forall a b. (a -> b) -> a -> b
$ Type -> Slice SubExp -> Slice SubExp
fullSlice Type
a_t [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
0::Int32)]
              bindArrayResult :: PatElemT attr -> SubExp -> m ()
bindArrayResult PatElemT attr
pe SubExp
se =
                [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [PatElemT attr -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT attr
pe] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$
                BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> BasicOp (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Type -> BasicOp (Lore m)
forall lore. [SubExp] -> Type -> BasicOp lore
ArrayLit [SubExp
se] (Type -> BasicOp (Lore m)) -> Type -> BasicOp (Lore m)
forall a b. (a -> b) -> a -> b
$ Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ PatElemT attr -> Type
forall attr. Typed attr => PatElemT attr -> Type
patElemType PatElemT attr
pe
              bindResult :: PatElemT attr -> SubExp -> m ()
bindResult PatElemT attr
pe SubExp
se =
                [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [PatElemT attr -> VName
forall attr. PatElemT attr -> VName
patElemName PatElemT attr
pe] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> BasicOp (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Lore m)
forall lore. SubExp -> BasicOp lore
SubExp SubExp
se
simplifyKnownIterationSOAC SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

data ArrayOp = ArrayIndexing Certificates VName (Slice SubExp)
             | ArrayRearrange Certificates VName [Int]
             | ArrayRotate Certificates VName [SubExp]
             | ArrayVar Certificates VName -- ^ Never constructed.
  deriving (ArrayOp -> ArrayOp -> Bool
(ArrayOp -> ArrayOp -> Bool)
-> (ArrayOp -> ArrayOp -> Bool) -> Eq ArrayOp
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ArrayOp -> ArrayOp -> Bool
$c/= :: ArrayOp -> ArrayOp -> Bool
== :: ArrayOp -> ArrayOp -> Bool
$c== :: ArrayOp -> ArrayOp -> Bool
Eq, Eq ArrayOp
Eq ArrayOp
-> (ArrayOp -> ArrayOp -> Ordering)
-> (ArrayOp -> ArrayOp -> Bool)
-> (ArrayOp -> ArrayOp -> Bool)
-> (ArrayOp -> ArrayOp -> Bool)
-> (ArrayOp -> ArrayOp -> Bool)
-> (ArrayOp -> ArrayOp -> ArrayOp)
-> (ArrayOp -> ArrayOp -> ArrayOp)
-> Ord ArrayOp
ArrayOp -> ArrayOp -> Bool
ArrayOp -> ArrayOp -> Ordering
ArrayOp -> ArrayOp -> ArrayOp
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ArrayOp -> ArrayOp -> ArrayOp
$cmin :: ArrayOp -> ArrayOp -> ArrayOp
max :: ArrayOp -> ArrayOp -> ArrayOp
$cmax :: ArrayOp -> ArrayOp -> ArrayOp
>= :: ArrayOp -> ArrayOp -> Bool
$c>= :: ArrayOp -> ArrayOp -> Bool
> :: ArrayOp -> ArrayOp -> Bool
$c> :: ArrayOp -> ArrayOp -> Bool
<= :: ArrayOp -> ArrayOp -> Bool
$c<= :: ArrayOp -> ArrayOp -> Bool
< :: ArrayOp -> ArrayOp -> Bool
$c< :: ArrayOp -> ArrayOp -> Bool
compare :: ArrayOp -> ArrayOp -> Ordering
$ccompare :: ArrayOp -> ArrayOp -> Ordering
$cp1Ord :: Eq ArrayOp
Ord, Int -> ArrayOp -> ShowS
[ArrayOp] -> ShowS
ArrayOp -> String
(Int -> ArrayOp -> ShowS)
-> (ArrayOp -> String) -> ([ArrayOp] -> ShowS) -> Show ArrayOp
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ArrayOp] -> ShowS
$cshowList :: [ArrayOp] -> ShowS
show :: ArrayOp -> String
$cshow :: ArrayOp -> String
showsPrec :: Int -> ArrayOp -> ShowS
$cshowsPrec :: Int -> ArrayOp -> ShowS
Show)

arrayOpArr :: ArrayOp -> VName
arrayOpArr :: ArrayOp -> VName
arrayOpArr (ArrayIndexing Certificates
_ VName
arr Slice SubExp
_) = VName
arr
arrayOpArr (ArrayRearrange Certificates
_ VName
arr [Int]
_) = VName
arr
arrayOpArr (ArrayRotate Certificates
_ VName
arr [SubExp]
_) = VName
arr
arrayOpArr (ArrayVar Certificates
_ VName
arr) = VName
arr

arrayOpCerts :: ArrayOp -> Certificates
arrayOpCerts :: ArrayOp -> Certificates
arrayOpCerts (ArrayIndexing Certificates
cs VName
_ Slice SubExp
_) = Certificates
cs
arrayOpCerts (ArrayRearrange Certificates
cs VName
_ [Int]
_) = Certificates
cs
arrayOpCerts (ArrayRotate Certificates
cs VName
_ [SubExp]
_) = Certificates
cs
arrayOpCerts (ArrayVar Certificates
cs VName
_) = Certificates
cs

isArrayOp :: Certificates -> AST.Exp (Wise SOACS) -> Maybe ArrayOp
isArrayOp :: Certificates -> Exp (Wise SOACS) -> Maybe ArrayOp
isArrayOp Certificates
cs (BasicOp (Index VName
arr Slice SubExp
slice)) =
  ArrayOp -> Maybe ArrayOp
forall a. a -> Maybe a
Just (ArrayOp -> Maybe ArrayOp) -> ArrayOp -> Maybe ArrayOp
forall a b. (a -> b) -> a -> b
$ Certificates -> VName -> Slice SubExp -> ArrayOp
ArrayIndexing Certificates
cs VName
arr Slice SubExp
slice
isArrayOp Certificates
cs (BasicOp (Rearrange [Int]
perm VName
arr)) =
  ArrayOp -> Maybe ArrayOp
forall a. a -> Maybe a
Just (ArrayOp -> Maybe ArrayOp) -> ArrayOp -> Maybe ArrayOp
forall a b. (a -> b) -> a -> b
$ Certificates -> VName -> [Int] -> ArrayOp
ArrayRearrange Certificates
cs VName
arr [Int]
perm
isArrayOp Certificates
cs (BasicOp (Rotate [SubExp]
rots VName
arr)) =
  ArrayOp -> Maybe ArrayOp
forall a. a -> Maybe a
Just (ArrayOp -> Maybe ArrayOp) -> ArrayOp -> Maybe ArrayOp
forall a b. (a -> b) -> a -> b
$ Certificates -> VName -> [SubExp] -> ArrayOp
ArrayRotate Certificates
cs VName
arr [SubExp]
rots
isArrayOp Certificates
_ Exp (Wise SOACS)
_ =
  Maybe ArrayOp
forall a. Maybe a
Nothing

fromArrayOp :: ArrayOp -> (Certificates, AST.Exp (Wise SOACS))
fromArrayOp :: ArrayOp -> (Certificates, Exp (Wise SOACS))
fromArrayOp (ArrayIndexing Certificates
cs VName
arr Slice SubExp
slice) = (Certificates
cs, BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp (Wise SOACS)
forall lore. VName -> Slice SubExp -> BasicOp lore
Index VName
arr Slice SubExp
slice)
fromArrayOp (ArrayRearrange Certificates
cs VName
arr [Int]
perm) = (Certificates
cs, BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp (Wise SOACS)
forall lore. [Int] -> VName -> BasicOp lore
Rearrange [Int]
perm VName
arr)
fromArrayOp (ArrayRotate Certificates
cs VName
arr [SubExp]
rots) = (Certificates
cs, BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> BasicOp (Wise SOACS)
forall lore. [SubExp] -> VName -> BasicOp lore
Rotate [SubExp]
rots VName
arr)
fromArrayOp (ArrayVar Certificates
cs VName
arr) = (Certificates
cs, BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Wise SOACS)
forall lore. SubExp -> BasicOp lore
SubExp (SubExp -> BasicOp (Wise SOACS)) -> SubExp -> BasicOp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
arr)

arrayOps :: AST.Body (Wise SOACS) -> S.Set (AST.Pattern (Wise SOACS), ArrayOp)
arrayOps :: Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
arrayOps = [Set (PatternT (VarWisdom, Type), ArrayOp)]
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall a. Monoid a => [a] -> a
mconcat ([Set (PatternT (VarWisdom, Type), ArrayOp)]
 -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> (Body (Wise SOACS)
    -> [Set (PatternT (VarWisdom, Type), ArrayOp)])
-> Body (Wise SOACS)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> [Stm (Wise SOACS)]
-> [Set (PatternT (VarWisdom, Type), ArrayOp)]
forall a b. (a -> b) -> [a] -> [b]
map Stm (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp)
onStm ([Stm (Wise SOACS)] -> [Set (PatternT (VarWisdom, Type), ArrayOp)])
-> (Body (Wise SOACS) -> [Stm (Wise SOACS)])
-> Body (Wise SOACS)
-> [Set (PatternT (VarWisdom, Type), ArrayOp)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms (Wise SOACS) -> [Stm (Wise SOACS)]
forall lore. Stms lore -> [Stm lore]
stmsToList (Stms (Wise SOACS) -> [Stm (Wise SOACS)])
-> (Body (Wise SOACS) -> Stms (Wise SOACS))
-> Body (Wise SOACS)
-> [Stm (Wise SOACS)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body (Wise SOACS) -> Stms (Wise SOACS)
forall lore. BodyT lore -> Stms lore
bodyStms
  where onStm :: Stm (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp)
onStm (Let Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
aux Exp (Wise SOACS)
e) =
          case Certificates -> Exp (Wise SOACS) -> Maybe ArrayOp
isArrayOp (StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux) Exp (Wise SOACS)
e of
            Just ArrayOp
op -> (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall a. a -> Set a
S.singleton (PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat, ArrayOp
op)
            Maybe ArrayOp
Nothing -> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
-> Set (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall s a. State s a -> s -> s
execState (Walker
  (Wise SOACS)
  (StateT (Set (PatternT (VarWisdom, Type), ArrayOp)) Identity)
-> Exp (Wise SOACS)
-> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
forall (m :: * -> *) lore.
Monad m =>
Walker lore m -> Exp lore -> m ()
walkExpM Walker
  (Wise SOACS)
  (StateT (Set (PatternT (VarWisdom, Type), ArrayOp)) Identity)
walker Exp (Wise SOACS)
e) Set (PatternT (VarWisdom, Type), ArrayOp)
forall a. Monoid a => a
mempty
        onOp :: SOAC (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp)
onOp = Writer
  (Set (PatternT (VarWisdom, Type), ArrayOp)) (SOAC (Wise SOACS))
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall w a. Writer w a -> w
execWriter (Writer
   (Set (PatternT (VarWisdom, Type), ArrayOp)) (SOAC (Wise SOACS))
 -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> (SOAC (Wise SOACS)
    -> Writer
         (Set (PatternT (VarWisdom, Type), ArrayOp)) (SOAC (Wise SOACS)))
-> SOAC (Wise SOACS)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOACMapper
  (Wise SOACS)
  (Wise SOACS)
  (WriterT (Set (PatternT (VarWisdom, Type), ArrayOp)) Identity)
-> SOAC (Wise SOACS)
-> Writer
     (Set (PatternT (VarWisdom, Type), ArrayOp)) (SOAC (Wise SOACS))
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore)
mapSOACM SOACMapper
  Any
  Any
  (WriterT (Set (PatternT (VarWisdom, Type), ArrayOp)) Identity)
forall (m :: * -> *) lore. Monad m => SOACMapper lore lore m
identitySOACMapper { mapOnSOACLambda :: Lambda (Wise SOACS)
-> WriterT
     (Set (PatternT (VarWisdom, Type), ArrayOp))
     Identity
     (Lambda (Wise SOACS))
mapOnSOACLambda = Lambda (Wise SOACS)
-> WriterT
     (Set (PatternT (VarWisdom, Type), ArrayOp))
     Identity
     (Lambda (Wise SOACS))
forall (m :: * -> *).
MonadWriter (Set (PatternT (VarWisdom, Type), ArrayOp)) m =>
Lambda (Wise SOACS) -> m (Lambda (Wise SOACS))
onLambda }
        onLambda :: Lambda (Wise SOACS) -> m (Lambda (Wise SOACS))
onLambda Lambda (Wise SOACS)
lam = do Set (PatternT (VarWisdom, Type), ArrayOp) -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Set (PatternT (VarWisdom, Type), ArrayOp) -> m ())
-> Set (PatternT (VarWisdom, Type), ArrayOp) -> m ()
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
arrayOps (Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp))
-> Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
lam
                          Lambda (Wise SOACS) -> m (Lambda (Wise SOACS))
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda (Wise SOACS)
lam
        walker :: Walker
  (Wise SOACS)
  (StateT (Set (PatternT (VarWisdom, Type), ArrayOp)) Identity)
walker = Walker
  (Wise SOACS)
  (StateT (Set (PatternT (VarWisdom, Type), ArrayOp)) Identity)
forall (m :: * -> *) lore. Monad m => Walker lore m
identityWalker { walkOnBody :: Body (Wise SOACS)
-> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
walkOnBody = (Set (PatternT (VarWisdom, Type), ArrayOp)
 -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set (PatternT (VarWisdom, Type), ArrayOp)
  -> Set (PatternT (VarWisdom, Type), ArrayOp))
 -> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ())
-> (Body (Wise SOACS)
    -> Set (PatternT (VarWisdom, Type), ArrayOp)
    -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> Body (Wise SOACS)
-> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall a. Semigroup a => a -> a -> a
(<>) (Set (PatternT (VarWisdom, Type), ArrayOp)
 -> Set (PatternT (VarWisdom, Type), ArrayOp)
 -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> (Body (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> Body (Wise SOACS)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp)
Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
arrayOps
                                , walkOnOp :: Op (Wise SOACS)
-> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
walkOnOp = (Set (PatternT (VarWisdom, Type), ArrayOp)
 -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set (PatternT (VarWisdom, Type), ArrayOp)
  -> Set (PatternT (VarWisdom, Type), ArrayOp))
 -> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ())
-> (SOAC (Wise SOACS)
    -> Set (PatternT (VarWisdom, Type), ArrayOp)
    -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> SOAC (Wise SOACS)
-> State (Set (PatternT (VarWisdom, Type), ArrayOp)) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall a. Semigroup a => a -> a -> a
(<>) (Set (PatternT (VarWisdom, Type), ArrayOp)
 -> Set (PatternT (VarWisdom, Type), ArrayOp)
 -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> (SOAC (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp))
-> SOAC (Wise SOACS)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
-> Set (PatternT (VarWisdom, Type), ArrayOp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC (Wise SOACS) -> Set (PatternT (VarWisdom, Type), ArrayOp)
onOp }

replaceArrayOps :: M.Map ArrayOp ArrayOp
                -> AST.Body (Wise SOACS) -> AST.Body (Wise SOACS)
replaceArrayOps :: Map ArrayOp ArrayOp -> Body (Wise SOACS) -> Body (Wise SOACS)
replaceArrayOps Map ArrayOp ArrayOp
substs (Body BodyAttr (Wise SOACS)
_ Stms (Wise SOACS)
stms [SubExp]
res) =
  Stms (Wise SOACS) -> [SubExp] -> Body (Wise SOACS)
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody ((Stm (Wise SOACS) -> Stm (Wise SOACS))
-> Stms (Wise SOACS) -> Stms (Wise SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm (Wise SOACS) -> Stm (Wise SOACS)
onStm Stms (Wise SOACS)
stms) [SubExp]
res
  where onStm :: Stm (Wise SOACS) -> Stm (Wise SOACS)
onStm (Let Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
aux Exp (Wise SOACS)
e) =
          let (Certificates
cs', Exp (Wise SOACS)
e') = Certificates
-> Exp (Wise SOACS) -> (Certificates, Exp (Wise SOACS))
onExp (StmAux (ExpWisdom, ()) -> Certificates
forall attr. StmAux attr -> Certificates
stmAuxCerts StmAux (ExpWisdom, ())
StmAux (ExpAttr (Wise SOACS))
aux) Exp (Wise SOACS)
e
          in Certificates -> Stm (Wise SOACS) -> Stm (Wise SOACS)
forall lore. Certificates -> Stm lore -> Stm lore
certify Certificates
cs' (Stm (Wise SOACS) -> Stm (Wise SOACS))
-> Stm (Wise SOACS) -> Stm (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [Ident] -> [Ident] -> Exp (Wise SOACS) -> Stm (Wise SOACS)
forall lore.
Bindable lore =>
[Ident] -> [Ident] -> Exp lore -> Stm lore
mkLet (PatternT (VarWisdom, Type) -> [Ident]
forall attr. Typed attr => PatternT attr -> [Ident]
patternContextIdents PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat) (PatternT (VarWisdom, Type) -> [Ident]
forall attr. Typed attr => PatternT attr -> [Ident]
patternValueIdents PatternT (VarWisdom, Type)
Pattern (Wise SOACS)
pat) Exp (Wise SOACS)
e'
        onExp :: Certificates
-> Exp (Wise SOACS) -> (Certificates, Exp (Wise SOACS))
onExp Certificates
cs Exp (Wise SOACS)
e
          | Just ArrayOp
op <- Certificates -> Exp (Wise SOACS) -> Maybe ArrayOp
isArrayOp Certificates
cs Exp (Wise SOACS)
e,
            Just ArrayOp
op' <- ArrayOp -> Map ArrayOp ArrayOp -> Maybe ArrayOp
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ArrayOp
op Map ArrayOp ArrayOp
substs =
              ArrayOp -> (Certificates, Exp (Wise SOACS))
fromArrayOp ArrayOp
op'
        onExp Certificates
cs Exp (Wise SOACS)
e = (Certificates
cs, Mapper (Wise SOACS) (Wise SOACS) Identity
-> Exp (Wise SOACS) -> Exp (Wise SOACS)
forall flore tlore.
Mapper flore tlore Identity -> Exp flore -> Exp tlore
mapExp Mapper (Wise SOACS) (Wise SOACS) Identity
mapper Exp (Wise SOACS)
e)
        mapper :: Mapper (Wise SOACS) (Wise SOACS) Identity
mapper = Mapper (Wise SOACS) (Wise SOACS) Identity
forall (m :: * -> *) lore. Monad m => Mapper lore lore m
identityMapper { mapOnBody :: Scope (Wise SOACS)
-> Body (Wise SOACS) -> Identity (Body (Wise SOACS))
mapOnBody = (Body (Wise SOACS) -> Identity (Body (Wise SOACS)))
-> Scope (Wise SOACS)
-> Body (Wise SOACS)
-> Identity (Body (Wise SOACS))
forall a b. a -> b -> a
const ((Body (Wise SOACS) -> Identity (Body (Wise SOACS)))
 -> Scope (Wise SOACS)
 -> Body (Wise SOACS)
 -> Identity (Body (Wise SOACS)))
-> (Body (Wise SOACS) -> Identity (Body (Wise SOACS)))
-> Scope (Wise SOACS)
-> Body (Wise SOACS)
-> Identity (Body (Wise SOACS))
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> Identity (Body (Wise SOACS))
forall (m :: * -> *) a. Monad m => a -> m a
return (Body (Wise SOACS) -> Identity (Body (Wise SOACS)))
-> (Body (Wise SOACS) -> Body (Wise SOACS))
-> Body (Wise SOACS)
-> Identity (Body (Wise SOACS))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map ArrayOp ArrayOp -> Body (Wise SOACS) -> Body (Wise SOACS)
replaceArrayOps Map ArrayOp ArrayOp
substs
                                , mapOnOp :: Op (Wise SOACS) -> Identity (Op (Wise SOACS))
mapOnOp = SOAC (Wise SOACS) -> Identity (SOAC (Wise SOACS))
forall (m :: * -> *) a. Monad m => a -> m a
return (SOAC (Wise SOACS) -> Identity (SOAC (Wise SOACS)))
-> (SOAC (Wise SOACS) -> SOAC (Wise SOACS))
-> SOAC (Wise SOACS)
-> Identity (SOAC (Wise SOACS))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC (Wise SOACS) -> SOAC (Wise SOACS)
onOp }
        onOp :: SOAC (Wise SOACS) -> SOAC (Wise SOACS)
onOp = Identity (SOAC (Wise SOACS)) -> SOAC (Wise SOACS)
forall a. Identity a -> a
runIdentity (Identity (SOAC (Wise SOACS)) -> SOAC (Wise SOACS))
-> (SOAC (Wise SOACS) -> Identity (SOAC (Wise SOACS)))
-> SOAC (Wise SOACS)
-> SOAC (Wise SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOACMapper (Wise SOACS) (Wise SOACS) Identity
-> SOAC (Wise SOACS) -> Identity (SOAC (Wise SOACS))
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore)
mapSOACM SOACMapper Any Any Identity
forall (m :: * -> *) lore. Monad m => SOACMapper lore lore m
identitySOACMapper { mapOnSOACLambda :: Lambda (Wise SOACS) -> Identity (Lambda (Wise SOACS))
mapOnSOACLambda = Lambda (Wise SOACS) -> Identity (Lambda (Wise SOACS))
forall (m :: * -> *) a. Monad m => a -> m a
return (Lambda (Wise SOACS) -> Identity (Lambda (Wise SOACS)))
-> (Lambda (Wise SOACS) -> Lambda (Wise SOACS))
-> Lambda (Wise SOACS)
-> Identity (Lambda (Wise SOACS))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda (Wise SOACS) -> Lambda (Wise SOACS)
onLambda }
        onLambda :: Lambda (Wise SOACS) -> Lambda (Wise SOACS)
onLambda Lambda (Wise SOACS)
lam = Lambda (Wise SOACS)
lam { lambdaBody :: Body (Wise SOACS)
lambdaBody = Map ArrayOp ArrayOp -> Body (Wise SOACS) -> Body (Wise SOACS)
replaceArrayOps Map ArrayOp ArrayOp
substs (Body (Wise SOACS) -> Body (Wise SOACS))
-> Body (Wise SOACS) -> Body (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
lam }

-- Turn
--
--    map (\i -> ... xs[i] ...) (iota n)
--
-- into
--
--    map (\i x -> ... x ...) (iota n) xs
--
-- This is not because we want to encourage the map-iota pattern, but
-- it may be present in generated code.  This is an unfortunately
-- expensive simplification rule, since it requires multiple passes
-- over the entire lambda body.  It only handles the very simplest
-- case - if you find yourself planning to extend it to handle more
-- complex situations (rotate or whatnot), consider turning it into a
-- separate compiler pass instead.
simplifyMapIota :: TopDownRuleOp (Wise SOACS)
simplifyMapIota :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
simplifyMapIota SymbolTable (Wise SOACS)
vtable Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma w (ScremaForm scan reduce map_lam) arrs)
  | Just (Param Type
p, VName
_) <- ((Param Type, VName) -> Bool)
-> [(Param Type, VName)] -> Maybe (Param Type, VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Param Type, VName) -> Bool
isIota ([Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
map_lam) [VName]
arrs),
    [ArrayOp]
indexings <- (ArrayOp -> Bool) -> [ArrayOp] -> [ArrayOp]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> ArrayOp -> Bool
indexesWith (Param Type -> VName
forall attr. Param attr -> VName
paramName Param Type
p)) ([ArrayOp] -> [ArrayOp]) -> [ArrayOp] -> [ArrayOp]
forall a b. (a -> b) -> a -> b
$ ((PatternT (VarWisdom, Type), ArrayOp) -> ArrayOp)
-> [(PatternT (VarWisdom, Type), ArrayOp)] -> [ArrayOp]
forall a b. (a -> b) -> [a] -> [b]
map (PatternT (VarWisdom, Type), ArrayOp) -> ArrayOp
forall a b. (a, b) -> b
snd ([(PatternT (VarWisdom, Type), ArrayOp)] -> [ArrayOp])
-> [(PatternT (VarWisdom, Type), ArrayOp)] -> [ArrayOp]
forall a b. (a -> b) -> a -> b
$ Set (PatternT (VarWisdom, Type), ArrayOp)
-> [(PatternT (VarWisdom, Type), ArrayOp)]
forall a. Set a -> [a]
S.toList (Set (PatternT (VarWisdom, Type), ArrayOp)
 -> [(PatternT (VarWisdom, Type), ArrayOp)])
-> Set (PatternT (VarWisdom, Type), ArrayOp)
-> [(PatternT (VarWisdom, Type), ArrayOp)]
forall a b. (a -> b) -> a -> b
$
                 Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
arrayOps (Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp))
-> Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [ArrayOp] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ArrayOp]
indexings = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      -- For each indexing with iota, add the corresponding array to
      -- the Screma, and construct a new lambda parameter.
      ([VName]
more_arrs, [Param Type]
more_params, [ArrayOp]
replacements) <-
        [(VName, Param Type, ArrayOp)]
-> ([VName], [Param Type], [ArrayOp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(VName, Param Type, ArrayOp)]
 -> ([VName], [Param Type], [ArrayOp]))
-> ([Maybe (VName, Param Type, ArrayOp)]
    -> [(VName, Param Type, ArrayOp)])
-> [Maybe (VName, Param Type, ArrayOp)]
-> ([VName], [Param Type], [ArrayOp])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (VName, Param Type, ArrayOp)]
-> [(VName, Param Type, ArrayOp)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (VName, Param Type, ArrayOp)]
 -> ([VName], [Param Type], [ArrayOp]))
-> RuleM (Wise SOACS) [Maybe (VName, Param Type, ArrayOp)]
-> RuleM (Wise SOACS) ([VName], [Param Type], [ArrayOp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ArrayOp
 -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp)))
-> [ArrayOp]
-> RuleM (Wise SOACS) [Maybe (VName, Param Type, ArrayOp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ArrayOp -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
mapOverArr [ArrayOp]
indexings
      let substs :: Map ArrayOp ArrayOp
substs = [(ArrayOp, ArrayOp)] -> Map ArrayOp ArrayOp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(ArrayOp, ArrayOp)] -> Map ArrayOp ArrayOp)
-> [(ArrayOp, ArrayOp)] -> Map ArrayOp ArrayOp
forall a b. (a -> b) -> a -> b
$ [ArrayOp] -> [ArrayOp] -> [(ArrayOp, ArrayOp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [ArrayOp]
indexings [ArrayOp]
replacements
          map_lam' :: Lambda (Wise SOACS)
map_lam' = Lambda (Wise SOACS)
map_lam { lambdaParams :: [LParam (Wise SOACS)]
lambdaParams = Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
map_lam [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
more_params
                             , lambdaBody :: Body (Wise SOACS)
lambdaBody = Map ArrayOp ArrayOp -> Body (Wise SOACS) -> Body (Wise SOACS)
replaceArrayOps Map ArrayOp ArrayOp
substs (Body (Wise SOACS) -> Body (Wise SOACS))
-> Body (Wise SOACS) -> Body (Wise SOACS)
forall a b. (a -> b) -> a -> b
$
                                            Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam
                             }
      Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w (Scan (Wise SOACS)
-> [Reduce (Wise SOACS)]
-> Lambda (Wise SOACS)
-> ScremaForm (Wise SOACS)
forall lore.
Scan lore -> [Reduce lore] -> Lambda lore -> ScremaForm lore
ScremaForm Scan (Wise SOACS)
scan [Reduce (Wise SOACS)]
reduce Lambda (Wise SOACS)
map_lam') ([VName]
arrs [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
more_arrs)
  where isIota :: (Param Type, VName) -> Bool
isIota (Param Type
_, VName
arr) = case VName
-> SymbolTable (Wise SOACS)
-> Maybe (BasicOp (Wise SOACS), Certificates)
forall lore.
VName -> SymbolTable lore -> Maybe (BasicOp lore, Certificates)
ST.lookupBasicOp VName
arr SymbolTable (Wise SOACS)
vtable of
                            Just (Iota SubExp
_ (Constant PrimValue
o) (Constant PrimValue
s) IntType
_, Certificates
_) ->
                              PrimValue -> Bool
zeroIsh PrimValue
o Bool -> Bool -> Bool
&& PrimValue -> Bool
oneIsh PrimValue
s
                            Maybe (BasicOp (Wise SOACS), Certificates)
_ -> Bool
False

        indexesWith :: VName -> ArrayOp -> Bool
indexesWith VName
v (ArrayIndexing Certificates
cs VName
arr (DimFix (Var VName
i) : Slice SubExp
_))
          | VName
arr VName -> SymbolTable (Wise SOACS) -> Bool
forall lore. VName -> SymbolTable lore -> Bool
`ST.elem` SymbolTable (Wise SOACS)
vtable,
            (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> SymbolTable (Wise SOACS) -> Bool
forall lore. VName -> SymbolTable lore -> Bool
`ST.elem` SymbolTable (Wise SOACS)
vtable) ([VName] -> Bool) -> [VName] -> Bool
forall a b. (a -> b) -> a -> b
$ Certificates -> [VName]
unCertificates Certificates
cs =
              VName
i VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v
        indexesWith VName
_ ArrayOp
_ = Bool
False

        mapOverArr :: ArrayOp -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
mapOverArr (ArrayIndexing Certificates
cs VName
arr Slice SubExp
slice) = do
          VName
arr_elem <- String -> RuleM (Wise SOACS) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> RuleM (Wise SOACS) VName)
-> String -> RuleM (Wise SOACS) VName
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString VName
arr String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_elem"
          Type
arr_t <- VName -> RuleM (Wise SOACS) Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
arr
          VName
arr' <- if Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 Type
arr_t SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
w
                  then VName -> RuleM (Wise SOACS) VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
arr
                  else Certificates
-> RuleM (Wise SOACS) VName -> RuleM (Wise SOACS) VName
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (RuleM (Wise SOACS) VName -> RuleM (Wise SOACS) VName)
-> RuleM (Wise SOACS) VName -> RuleM (Wise SOACS) VName
forall a b. (a -> b) -> a -> b
$ String
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp (VName -> String
baseString VName
arr String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_prefix") (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) VName)
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) VName
forall a b. (a -> b) -> a -> b
$
                       BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp (Wise SOACS)
forall lore. VName -> Slice SubExp -> BasicOp lore
Index VName
arr (Slice SubExp -> BasicOp (Wise SOACS))
-> Slice SubExp -> BasicOp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$
                       Type -> Slice SubExp -> Slice SubExp
fullSlice Type
arr_t [SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
0) SubExp
w (IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
1)]
          Maybe (VName, Param Type, ArrayOp)
-> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (VName, Param Type, ArrayOp)
 -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp)))
-> Maybe (VName, Param Type, ArrayOp)
-> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
forall a b. (a -> b) -> a -> b
$ (VName, Param Type, ArrayOp) -> Maybe (VName, Param Type, ArrayOp)
forall a. a -> Maybe a
Just (VName
arr',
                         VName -> Type -> Param Type
forall attr. VName -> attr -> Param attr
Param VName
arr_elem (Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType Type
arr_t),
                         Certificates -> VName -> Slice SubExp -> ArrayOp
ArrayIndexing Certificates
cs VName
arr_elem (Int -> Slice SubExp -> Slice SubExp
forall a. Int -> [a] -> [a]
drop Int
1 Slice SubExp
slice))

        mapOverArr ArrayOp
_ = Maybe (VName, Param Type, ArrayOp)
-> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (VName, Param Type, ArrayOp)
forall a. Maybe a
Nothing

simplifyMapIota  SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ = Rule (Wise SOACS)
forall lore. Rule lore
Skip

-- If a Screma's map function contains a transformation
-- (e.g. transpose) on a parameter, create a new parameter
-- corresponding to that transformation performed on the rows of the
-- full array.
moveTransformToInput :: TopDownRuleOp (Wise SOACS)
moveTransformToInput :: RuleOp (Wise SOACS) (SymbolTable (Wise SOACS))
moveTransformToInput SymbolTable (Wise SOACS)
vtable Pattern (Wise SOACS)
pat StmAux (ExpAttr (Wise SOACS))
_ (Screma w (ScremaForm scan reduce map_lam) arrs)
  | [ArrayOp]
ops <- ((PatternT (VarWisdom, Type), ArrayOp) -> ArrayOp)
-> [(PatternT (VarWisdom, Type), ArrayOp)] -> [ArrayOp]
forall a b. (a -> b) -> [a] -> [b]
map (PatternT (VarWisdom, Type), ArrayOp) -> ArrayOp
forall a b. (a, b) -> b
snd ([(PatternT (VarWisdom, Type), ArrayOp)] -> [ArrayOp])
-> [(PatternT (VarWisdom, Type), ArrayOp)] -> [ArrayOp]
forall a b. (a -> b) -> a -> b
$ ((PatternT (VarWisdom, Type), ArrayOp) -> Bool)
-> [(PatternT (VarWisdom, Type), ArrayOp)]
-> [(PatternT (VarWisdom, Type), ArrayOp)]
forall a. (a -> Bool) -> [a] -> [a]
filter (PatternT (VarWisdom, Type), ArrayOp) -> Bool
arrayIsMapParam ([(PatternT (VarWisdom, Type), ArrayOp)]
 -> [(PatternT (VarWisdom, Type), ArrayOp)])
-> [(PatternT (VarWisdom, Type), ArrayOp)]
-> [(PatternT (VarWisdom, Type), ArrayOp)]
forall a b. (a -> b) -> a -> b
$ Set (PatternT (VarWisdom, Type), ArrayOp)
-> [(PatternT (VarWisdom, Type), ArrayOp)]
forall a. Set a -> [a]
S.toList (Set (PatternT (VarWisdom, Type), ArrayOp)
 -> [(PatternT (VarWisdom, Type), ArrayOp)])
-> Set (PatternT (VarWisdom, Type), ArrayOp)
-> [(PatternT (VarWisdom, Type), ArrayOp)]
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
arrayOps (Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp))
-> Body (Wise SOACS) -> Set (Pattern (Wise SOACS), ArrayOp)
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [ArrayOp] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ArrayOp]
ops = RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall lore. RuleM lore () -> Rule lore
Simplify (RuleM (Wise SOACS) () -> Rule (Wise SOACS))
-> RuleM (Wise SOACS) () -> Rule (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ do
      ([VName]
more_arrs, [Param Type]
more_params, [ArrayOp]
replacements) <-
        [(VName, Param Type, ArrayOp)]
-> ([VName], [Param Type], [ArrayOp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(VName, Param Type, ArrayOp)]
 -> ([VName], [Param Type], [ArrayOp]))
-> ([Maybe (VName, Param Type, ArrayOp)]
    -> [(VName, Param Type, ArrayOp)])
-> [Maybe (VName, Param Type, ArrayOp)]
-> ([VName], [Param Type], [ArrayOp])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (VName, Param Type, ArrayOp)]
-> [(VName, Param Type, ArrayOp)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (VName, Param Type, ArrayOp)]
 -> ([VName], [Param Type], [ArrayOp]))
-> RuleM (Wise SOACS) [Maybe (VName, Param Type, ArrayOp)]
-> RuleM (Wise SOACS) ([VName], [Param Type], [ArrayOp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ArrayOp
 -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp)))
-> [ArrayOp]
-> RuleM (Wise SOACS) [Maybe (VName, Param Type, ArrayOp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ArrayOp -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
mapOverArr [ArrayOp]
ops

      Bool -> RuleM (Wise SOACS) () -> RuleM (Wise SOACS) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
more_arrs) RuleM (Wise SOACS) ()
forall lore a. RuleM lore a
cannotSimplify

      let substs :: Map ArrayOp ArrayOp
substs = [(ArrayOp, ArrayOp)] -> Map ArrayOp ArrayOp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(ArrayOp, ArrayOp)] -> Map ArrayOp ArrayOp)
-> [(ArrayOp, ArrayOp)] -> Map ArrayOp ArrayOp
forall a b. (a -> b) -> a -> b
$ [ArrayOp] -> [ArrayOp] -> [(ArrayOp, ArrayOp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [ArrayOp]
ops [ArrayOp]
replacements
          map_lam' :: Lambda (Wise SOACS)
map_lam' = Lambda (Wise SOACS)
map_lam { lambdaParams :: [LParam (Wise SOACS)]
lambdaParams = Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
map_lam [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
more_params
                             , lambdaBody :: Body (Wise SOACS)
lambdaBody = Map ArrayOp ArrayOp -> Body (Wise SOACS) -> Body (Wise SOACS)
replaceArrayOps Map ArrayOp ArrayOp
substs (Body (Wise SOACS) -> Body (Wise SOACS))
-> Body (Wise SOACS) -> Body (Wise SOACS)
forall a b. (a -> b) -> a -> b
$
                                            Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam
                             }

      Pattern (Lore (RuleM (Wise SOACS)))
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (RuleM (Wise SOACS)))
Pattern (Wise SOACS)
pat (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ())
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise SOACS) -> Exp (Wise SOACS)
forall lore. Op lore -> ExpT lore
Op (Op (Wise SOACS) -> Exp (Wise SOACS))
-> Op (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Wise SOACS) -> [VName] -> SOAC (Wise SOACS)
forall lore. SubExp -> ScremaForm lore -> [VName] -> SOAC lore
Screma SubExp
w (Scan (Wise SOACS)
-> [Reduce (Wise SOACS)]
-> Lambda (Wise SOACS)
-> ScremaForm (Wise SOACS)
forall lore.
Scan lore -> [Reduce lore] -> Lambda lore -> ScremaForm lore
ScremaForm Scan (Wise SOACS)
scan [Reduce (Wise SOACS)]
reduce Lambda (Wise SOACS)
map_lam') ([VName]
arrs [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
more_arrs)

  where map_param_names :: [VName]
map_param_names = (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall attr. Param attr -> VName
paramName (Lambda (Wise SOACS) -> [LParam (Wise SOACS)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda (Wise SOACS)
map_lam)
        topLevelPattern :: PatternT (VarWisdom, Type) -> Bool
topLevelPattern = (PatternT (VarWisdom, Type)
-> Seq (PatternT (VarWisdom, Type)) -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` (Stm (Wise SOACS) -> PatternT (VarWisdom, Type))
-> Stms (Wise SOACS) -> Seq (PatternT (VarWisdom, Type))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm (Wise SOACS) -> PatternT (VarWisdom, Type)
forall lore. Stm lore -> Pattern lore
stmPattern (Body (Wise SOACS) -> Stms (Wise SOACS)
forall lore. BodyT lore -> Stms lore
bodyStms (Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam)))
        onlyUsedOnce :: VName -> Bool
onlyUsedOnce VName
arr =
          case (Stm (Wise SOACS) -> Bool)
-> [Stm (Wise SOACS)] -> [Stm (Wise SOACS)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName
arr VName -> Names -> Bool
`nameIn`) (Names -> Bool)
-> (Stm (Wise SOACS) -> Names) -> Stm (Wise SOACS) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm (Wise SOACS) -> Names
forall a. FreeIn a => a -> Names
freeIn) ([Stm (Wise SOACS)] -> [Stm (Wise SOACS)])
-> [Stm (Wise SOACS)] -> [Stm (Wise SOACS)]
forall a b. (a -> b) -> a -> b
$ Stms (Wise SOACS) -> [Stm (Wise SOACS)]
forall lore. Stms lore -> [Stm lore]
stmsToList (Stms (Wise SOACS) -> [Stm (Wise SOACS)])
-> Stms (Wise SOACS) -> [Stm (Wise SOACS)]
forall a b. (a -> b) -> a -> b
$ Body (Wise SOACS) -> Stms (Wise SOACS)
forall lore. BodyT lore -> Stms lore
bodyStms (Body (Wise SOACS) -> Stms (Wise SOACS))
-> Body (Wise SOACS) -> Stms (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda (Wise SOACS) -> Body (Wise SOACS)
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda (Wise SOACS)
map_lam of
            Stm (Wise SOACS)
_ : Stm (Wise SOACS)
_ : [Stm (Wise SOACS)]
_ -> Bool
False
            [Stm (Wise SOACS)]
_ -> Bool
True

        -- It's not just about whether the array is a parameter;
        -- everything else must be map-invariant.
        arrayIsMapParam :: (PatternT (VarWisdom, Type), ArrayOp) -> Bool
arrayIsMapParam (PatternT (VarWisdom, Type)
pat', ArrayIndexing Certificates
cs VName
arr Slice SubExp
slice) =
          VName
arr VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
map_param_names Bool -> Bool -> Bool
&&
          (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> SymbolTable (Wise SOACS) -> Bool
forall lore. VName -> SymbolTable lore -> Bool
`ST.elem` SymbolTable (Wise SOACS)
vtable) (Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Certificates -> Names
forall a. FreeIn a => a -> Names
freeIn Certificates
cs Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Slice SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn Slice SubExp
slice) Bool -> Bool -> Bool
&&
          Bool -> Bool
not (Slice SubExp -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Slice SubExp
slice) Bool -> Bool -> Bool
&&
          (Bool -> Bool
not ([SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([SubExp] -> Bool) -> [SubExp] -> Bool
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice) Bool -> Bool -> Bool
|| (PatternT (VarWisdom, Type) -> Bool
topLevelPattern PatternT (VarWisdom, Type)
pat' Bool -> Bool -> Bool
&& VName -> Bool
onlyUsedOnce VName
arr))
        arrayIsMapParam (PatternT (VarWisdom, Type)
_, ArrayRearrange Certificates
cs VName
arr [Int]
perm) =
          VName
arr VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
map_param_names Bool -> Bool -> Bool
&&
          (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> SymbolTable (Wise SOACS) -> Bool
forall lore. VName -> SymbolTable lore -> Bool
`ST.elem` SymbolTable (Wise SOACS)
vtable) (Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Certificates -> Names
forall a. FreeIn a => a -> Names
freeIn Certificates
cs) Bool -> Bool -> Bool
&&
          Bool -> Bool
not ([Int] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
perm)
        arrayIsMapParam (PatternT (VarWisdom, Type)
_, ArrayRotate Certificates
cs VName
arr [SubExp]
rots) =
          VName
arr VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
map_param_names Bool -> Bool -> Bool
&&
          (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> SymbolTable (Wise SOACS) -> Bool
forall lore. VName -> SymbolTable lore -> Bool
`ST.elem` SymbolTable (Wise SOACS)
vtable) (Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Certificates -> Names
forall a. FreeIn a => a -> Names
freeIn Certificates
cs Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
rots)
        arrayIsMapParam (PatternT (VarWisdom, Type)
_, ArrayVar{}) =
          Bool
False

        mapOverArr :: ArrayOp -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
mapOverArr ArrayOp
op
         | Just (VName
_, VName
arr) <- ((VName, VName) -> Bool)
-> [(VName, VName)] -> Maybe (VName, VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
==ArrayOp -> VName
arrayOpArr ArrayOp
op) (VName -> Bool)
-> ((VName, VName) -> VName) -> (VName, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, VName) -> VName
forall a b. (a, b) -> a
fst) ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
map_param_names [VName]
arrs) = do
             Type
arr_t <- VName -> RuleM (Wise SOACS) Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
arr
             let whole_dim :: DimIndex SubExp
whole_dim = SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
0) (Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 Type
arr_t) (IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
1)
             VName
arr_transformed <- Certificates
-> RuleM (Wise SOACS) VName -> RuleM (Wise SOACS) VName
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (ArrayOp -> Certificates
arrayOpCerts ArrayOp
op) (RuleM (Wise SOACS) VName -> RuleM (Wise SOACS) VName)
-> RuleM (Wise SOACS) VName -> RuleM (Wise SOACS) VName
forall a b. (a -> b) -> a -> b
$
                                String
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp (VName -> String
baseString VName
arr String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_transformed") (Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) VName)
-> Exp (Lore (RuleM (Wise SOACS))) -> RuleM (Wise SOACS) VName
forall a b. (a -> b) -> a -> b
$
                                case ArrayOp
op of
                                  ArrayIndexing Certificates
_ VName
_ Slice SubExp
slice ->
                                    BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp (Wise SOACS)
forall lore. VName -> Slice SubExp -> BasicOp lore
Index VName
arr (Slice SubExp -> BasicOp (Wise SOACS))
-> Slice SubExp -> BasicOp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ DimIndex SubExp
whole_dim DimIndex SubExp -> Slice SubExp -> Slice SubExp
forall a. a -> [a] -> [a]
: Slice SubExp
slice
                                  ArrayRearrange Certificates
_ VName
_ [Int]
perm ->
                                    BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp (Wise SOACS)
forall lore. [Int] -> VName -> BasicOp lore
Rearrange (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Int]
perm) VName
arr
                                  ArrayRotate Certificates
_ VName
_ [SubExp]
rots ->
                                    BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> BasicOp (Wise SOACS)
forall lore. [SubExp] -> VName -> BasicOp lore
Rotate (IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
0 SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
rots) VName
arr
                                  ArrayVar{} ->
                                    BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Wise SOACS) -> Exp (Wise SOACS))
-> BasicOp (Wise SOACS) -> Exp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Wise SOACS)
forall lore. SubExp -> BasicOp lore
SubExp (SubExp -> BasicOp (Wise SOACS)) -> SubExp -> BasicOp (Wise SOACS)
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
arr
             Type
arr_transformed_t <- VName -> RuleM (Wise SOACS) Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
arr_transformed
             VName
arr_transformed_row <- String -> RuleM (Wise SOACS) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> RuleM (Wise SOACS) VName)
-> String -> RuleM (Wise SOACS) VName
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString VName
arr String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_transformed_row"
             Maybe (VName, Param Type, ArrayOp)
-> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (VName, Param Type, ArrayOp)
 -> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp)))
-> Maybe (VName, Param Type, ArrayOp)
-> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
forall a b. (a -> b) -> a -> b
$ (VName, Param Type, ArrayOp) -> Maybe (VName, Param Type, ArrayOp)
forall a. a -> Maybe a
Just (VName
arr_transformed,
                            VName -> Type -> Param Type
forall attr. VName -> attr -> Param attr
Param VName
arr_transformed_row (Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType Type
arr_transformed_t),
                            Certificates -> VName -> ArrayOp
ArrayVar Certificates
forall a. Monoid a => a
mempty VName
arr_transformed_row)

        mapOverArr ArrayOp
_ = Maybe (VName, Param Type, ArrayOp)
-> RuleM (Wise SOACS) (Maybe (VName, Param Type, ArrayOp))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (VName, Param Type, ArrayOp)
forall a. Maybe a
Nothing

moveTransformToInput SymbolTable (Wise SOACS)
_ Pattern (Wise SOACS)
_ StmAux (ExpAttr (Wise SOACS))
_ Op (Wise SOACS)
_ =
  Rule (Wise SOACS)
forall lore. Rule lore
Skip