{-# LANGUAGE TypeFamilies #-}

-- | Facilities for fusing two SOACs.
--
-- When the fusion algorithm decides that it's worth fusing two SOAC
-- statements, this is the module that tries to see if that's
-- possible.  May involve massaging either producer or consumer in
-- various ways.
module Futhark.Optimise.Fusion.TryFusion
  ( FusedSOAC (..),
    Mode (..),
    attemptFusion,
  )
where

import Control.Applicative
import Control.Arrow (first)
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Data.Either (partitionEithers)
import Data.List (find)
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.HORep.MapNest (MapNest)
import Futhark.Analysis.HORep.MapNest qualified as MapNest
import Futhark.Analysis.HORep.SOAC qualified as SOAC
import Futhark.Construct
import Futhark.IR.SOACS hiding (SOAC (..))
import Futhark.IR.SOACS qualified as Futhark
import Futhark.Optimise.Fusion.Composing
import Futhark.Optimise.Fusion.Screma
import Futhark.Transform.ISRWIM (rwimPossible)
import Futhark.Transform.Rename (renameLambda)
import Futhark.Transform.Substitute

newtype TryFusion a
  = TryFusion
      ( ReaderT
          (Scope SOACS)
          (StateT VNameSource Maybe)
          a
      )
  deriving
    ( (forall a b. (a -> b) -> TryFusion a -> TryFusion b)
-> (forall a b. a -> TryFusion b -> TryFusion a)
-> Functor TryFusion
forall a b. a -> TryFusion b -> TryFusion a
forall a b. (a -> b) -> TryFusion a -> TryFusion b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> TryFusion a -> TryFusion b
fmap :: forall a b. (a -> b) -> TryFusion a -> TryFusion b
$c<$ :: forall a b. a -> TryFusion b -> TryFusion a
<$ :: forall a b. a -> TryFusion b -> TryFusion a
Functor,
      Functor TryFusion
Functor TryFusion =>
(forall a. a -> TryFusion a)
-> (forall a b. TryFusion (a -> b) -> TryFusion a -> TryFusion b)
-> (forall a b c.
    (a -> b -> c) -> TryFusion a -> TryFusion b -> TryFusion c)
-> (forall a b. TryFusion a -> TryFusion b -> TryFusion b)
-> (forall a b. TryFusion a -> TryFusion b -> TryFusion a)
-> Applicative TryFusion
forall a. a -> TryFusion a
forall a b. TryFusion a -> TryFusion b -> TryFusion a
forall a b. TryFusion a -> TryFusion b -> TryFusion b
forall a b. TryFusion (a -> b) -> TryFusion a -> TryFusion b
forall a b c.
(a -> b -> c) -> TryFusion a -> TryFusion b -> TryFusion c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> TryFusion a
pure :: forall a. a -> TryFusion a
$c<*> :: forall a b. TryFusion (a -> b) -> TryFusion a -> TryFusion b
<*> :: forall a b. TryFusion (a -> b) -> TryFusion a -> TryFusion b
$cliftA2 :: forall a b c.
(a -> b -> c) -> TryFusion a -> TryFusion b -> TryFusion c
liftA2 :: forall a b c.
(a -> b -> c) -> TryFusion a -> TryFusion b -> TryFusion c
$c*> :: forall a b. TryFusion a -> TryFusion b -> TryFusion b
*> :: forall a b. TryFusion a -> TryFusion b -> TryFusion b
$c<* :: forall a b. TryFusion a -> TryFusion b -> TryFusion a
<* :: forall a b. TryFusion a -> TryFusion b -> TryFusion a
Applicative,
      Applicative TryFusion
Applicative TryFusion =>
(forall a. TryFusion a)
-> (forall a. TryFusion a -> TryFusion a -> TryFusion a)
-> (forall a. TryFusion a -> TryFusion [a])
-> (forall a. TryFusion a -> TryFusion [a])
-> Alternative TryFusion
forall a. TryFusion a
forall a. TryFusion a -> TryFusion [a]
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *).
Applicative f =>
(forall a. f a)
-> (forall a. f a -> f a -> f a)
-> (forall a. f a -> f [a])
-> (forall a. f a -> f [a])
-> Alternative f
$cempty :: forall a. TryFusion a
empty :: forall a. TryFusion a
$c<|> :: forall a. TryFusion a -> TryFusion a -> TryFusion a
<|> :: forall a. TryFusion a -> TryFusion a -> TryFusion a
$csome :: forall a. TryFusion a -> TryFusion [a]
some :: forall a. TryFusion a -> TryFusion [a]
$cmany :: forall a. TryFusion a -> TryFusion [a]
many :: forall a. TryFusion a -> TryFusion [a]
Alternative,
      Applicative TryFusion
Applicative TryFusion =>
(forall a b. TryFusion a -> (a -> TryFusion b) -> TryFusion b)
-> (forall a b. TryFusion a -> TryFusion b -> TryFusion b)
-> (forall a. a -> TryFusion a)
-> Monad TryFusion
forall a. a -> TryFusion a
forall a b. TryFusion a -> TryFusion b -> TryFusion b
forall a b. TryFusion a -> (a -> TryFusion b) -> TryFusion b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. TryFusion a -> (a -> TryFusion b) -> TryFusion b
>>= :: forall a b. TryFusion a -> (a -> TryFusion b) -> TryFusion b
$c>> :: forall a b. TryFusion a -> TryFusion b -> TryFusion b
>> :: forall a b. TryFusion a -> TryFusion b -> TryFusion b
$creturn :: forall a. a -> TryFusion a
return :: forall a. a -> TryFusion a
Monad,
      Monad TryFusion
Monad TryFusion =>
(forall a. String -> TryFusion a) -> MonadFail TryFusion
forall a. String -> TryFusion a
forall (m :: * -> *).
Monad m =>
(forall a. String -> m a) -> MonadFail m
$cfail :: forall a. String -> TryFusion a
fail :: forall a. String -> TryFusion a
MonadFail,
      Monad TryFusion
TryFusion VNameSource
Monad TryFusion =>
TryFusion VNameSource
-> (VNameSource -> TryFusion ()) -> MonadFreshNames TryFusion
VNameSource -> TryFusion ()
forall (m :: * -> *).
Monad m =>
m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
$cgetNameSource :: TryFusion VNameSource
getNameSource :: TryFusion VNameSource
$cputNameSource :: VNameSource -> TryFusion ()
putNameSource :: VNameSource -> TryFusion ()
MonadFreshNames,
      HasScope SOACS,
      LocalScope SOACS
    )

tryFusion ::
  (MonadFreshNames m) =>
  TryFusion a ->
  Scope SOACS ->
  m (Maybe a)
tryFusion :: forall (m :: * -> *) a.
MonadFreshNames m =>
TryFusion a -> Scope SOACS -> m (Maybe a)
tryFusion (TryFusion ReaderT (Scope SOACS) (StateT VNameSource Maybe) a
m) Scope SOACS
types = (VNameSource -> (Maybe a, VNameSource)) -> m (Maybe a)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (Maybe a, VNameSource)) -> m (Maybe a))
-> (VNameSource -> (Maybe a, VNameSource)) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
  case StateT VNameSource Maybe a -> VNameSource -> Maybe (a, VNameSource)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (ReaderT (Scope SOACS) (StateT VNameSource Maybe) a
-> Scope SOACS -> StateT VNameSource Maybe a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Scope SOACS) (StateT VNameSource Maybe) a
m Scope SOACS
types) VNameSource
src of
    Just (a
x, VNameSource
src') -> (a -> Maybe a
forall a. a -> Maybe a
Just a
x, VNameSource
src')
    Maybe (a, VNameSource)
Nothing -> (Maybe a
forall a. Maybe a
Nothing, VNameSource
src)

liftMaybe :: Maybe a -> TryFusion a
liftMaybe :: forall a. Maybe a -> TryFusion a
liftMaybe Maybe a
Nothing = String -> TryFusion a
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Nothing"
liftMaybe (Just a
x) = a -> TryFusion a
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

type SOAC = SOAC.SOAC SOACS

inputToOutput :: SOAC.Input -> Maybe (SOAC.ArrayTransform, SOAC.Input)
inputToOutput :: Input -> Maybe (ArrayTransform, Input)
inputToOutput (SOAC.Input ArrayTransforms
ts VName
ia Type
iat) =
  case ArrayTransforms -> ViewF
SOAC.viewf ArrayTransforms
ts of
    ArrayTransform
t SOAC.:< ArrayTransforms
ts' -> (ArrayTransform, Input) -> Maybe (ArrayTransform, Input)
forall a. a -> Maybe a
Just (ArrayTransform
t, ArrayTransforms -> VName -> Type -> Input
SOAC.Input ArrayTransforms
ts' VName
ia Type
iat)
    ViewF
SOAC.EmptyF -> Maybe (ArrayTransform, Input)
forall a. Maybe a
Nothing

-- | A fused SOAC contains a bit of extra information.
data FusedSOAC = FusedSOAC
  { -- | The actual SOAC.
    FusedSOAC -> SOAC SOACS
fsSOAC :: SOAC,
    -- | A transformation to be applied to *all* results of the SOAC.
    FusedSOAC -> ArrayTransforms
fsOutputTransform :: SOAC.ArrayTransforms,
    -- | The outputs of the SOAC (i.e. the names in the pattern that
    -- the result of this SOAC should be bound to).
    FusedSOAC -> [VName]
fsOutNames :: [VName]
  }
  deriving (Int -> FusedSOAC -> ShowS
[FusedSOAC] -> ShowS
FusedSOAC -> String
(Int -> FusedSOAC -> ShowS)
-> (FusedSOAC -> String)
-> ([FusedSOAC] -> ShowS)
-> Show FusedSOAC
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> FusedSOAC -> ShowS
showsPrec :: Int -> FusedSOAC -> ShowS
$cshow :: FusedSOAC -> String
show :: FusedSOAC -> String
$cshowList :: [FusedSOAC] -> ShowS
showList :: [FusedSOAC] -> ShowS
Show)

inputs :: FusedSOAC -> [SOAC.Input]
inputs :: FusedSOAC -> [Input]
inputs = SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
SOAC.inputs (SOAC SOACS -> [Input])
-> (FusedSOAC -> SOAC SOACS) -> FusedSOAC -> [Input]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FusedSOAC -> SOAC SOACS
fsSOAC

setInputs :: [SOAC.Input] -> FusedSOAC -> FusedSOAC
setInputs :: [Input] -> FusedSOAC -> FusedSOAC
setInputs [Input]
inps FusedSOAC
ker = FusedSOAC
ker {fsSOAC = inps `SOAC.setInputs` fsSOAC ker}

tryOptimizeSOAC ::
  Mode ->
  Names ->
  [VName] ->
  SOAC ->
  FusedSOAC ->
  TryFusion FusedSOAC
tryOptimizeSOAC :: Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
tryOptimizeSOAC Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker = do
  (soac', ots) <- Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
optimizeSOAC Maybe [VName]
forall a. Maybe a
Nothing SOAC SOACS
soac ArrayTransforms
forall a. Monoid a => a
mempty
  let ker' = (Input -> Input) -> [Input] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (ArrayTransforms -> Input -> Input
addInitialTransformIfRelevant ArrayTransforms
ots) (FusedSOAC -> [Input]
inputs FusedSOAC
ker) [Input] -> FusedSOAC -> FusedSOAC
`setInputs` FusedSOAC
ker
      outIdents = (VName -> Type -> Ident) -> [VName] -> [Type] -> [Ident]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Ident
Ident [VName]
outVars ([Type] -> [Ident]) -> [Type] -> [Ident]
forall a b. (a -> b) -> a -> b
$ SOAC SOACS -> [Type]
forall rep. SOAC rep -> [Type]
SOAC.typeOf SOAC SOACS
soac'
      ker'' = [Ident] -> FusedSOAC -> FusedSOAC
fixInputTypes [Ident]
outIdents FusedSOAC
ker'
  applyFusionRules mode unfus_nms outVars soac' ker''
  where
    addInitialTransformIfRelevant :: ArrayTransforms -> Input -> Input
addInitialTransformIfRelevant ArrayTransforms
ots Input
inp
      | Input -> VName
SOAC.inputArray Input
inp VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
outVars =
          ArrayTransforms -> Input -> Input
SOAC.addInitialTransforms ArrayTransforms
ots Input
inp
      | Bool
otherwise =
          Input
inp

tryOptimizeKernel ::
  Mode ->
  Names ->
  [VName] ->
  SOAC ->
  FusedSOAC ->
  TryFusion FusedSOAC
tryOptimizeKernel :: Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
tryOptimizeKernel Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker = do
  ker' <- Maybe [VName] -> FusedSOAC -> TryFusion FusedSOAC
optimizeKernel ([VName] -> Maybe [VName]
forall a. a -> Maybe a
Just [VName]
outVars) FusedSOAC
ker
  applyFusionRules mode unfus_nms outVars soac ker'

tryExposeInputs ::
  Mode ->
  Names ->
  [VName] ->
  SOAC ->
  FusedSOAC ->
  TryFusion FusedSOAC
tryExposeInputs :: Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
tryExposeInputs Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker = do
  (ker', ots) <- [VName] -> FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs [VName]
outVars FusedSOAC
ker
  if SOAC.nullTransforms ots
    then fuseSOACwithKer mode unfus_nms outVars soac ker'
    else do
      guard $ unfus_nms == mempty
      (soac', ots') <- pullOutputTransforms soac ots
      let outIdents = (VName -> Type -> Ident) -> [VName] -> [Type] -> [Ident]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Ident
Ident [VName]
outVars ([Type] -> [Ident]) -> [Type] -> [Ident]
forall a b. (a -> b) -> a -> b
$ SOAC SOACS -> [Type]
forall rep. SOAC rep -> [Type]
SOAC.typeOf SOAC SOACS
soac'
          ker'' = [Ident] -> FusedSOAC -> FusedSOAC
fixInputTypes [Ident]
outIdents FusedSOAC
ker'
      if SOAC.nullTransforms ots'
        then applyFusionRules mode unfus_nms outVars soac' ker''
        else fail "tryExposeInputs could not pull SOAC transforms"

fixInputTypes :: [Ident] -> FusedSOAC -> FusedSOAC
fixInputTypes :: [Ident] -> FusedSOAC -> FusedSOAC
fixInputTypes [Ident]
outIdents FusedSOAC
ker =
  FusedSOAC
ker {fsSOAC = fixInputTypes' $ fsSOAC ker}
  where
    fixInputTypes' :: SOAC SOACS -> SOAC SOACS
fixInputTypes' SOAC SOACS
soac =
      (Input -> Input) -> [Input] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map Input -> Input
fixInputType (SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
SOAC.inputs SOAC SOACS
soac) [Input] -> SOAC SOACS -> SOAC SOACS
forall rep. [Input] -> SOAC rep -> SOAC rep
`SOAC.setInputs` SOAC SOACS
soac
    fixInputType :: Input -> Input
fixInputType (SOAC.Input ArrayTransforms
ts VName
v Type
_)
      | Just Ident
v' <- (Ident -> Bool) -> [Ident] -> Maybe Ident
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool) -> (Ident -> VName) -> Ident -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> VName
identName) [Ident]
outIdents =
          ArrayTransforms -> VName -> Type -> Input
SOAC.Input ArrayTransforms
ts VName
v (Type -> Input) -> Type -> Input
forall a b. (a -> b) -> a -> b
$ Ident -> Type
identType Ident
v'
    fixInputType Input
inp = Input
inp

applyFusionRules ::
  Mode ->
  Names ->
  [VName] ->
  SOAC ->
  FusedSOAC ->
  TryFusion FusedSOAC
applyFusionRules :: Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
applyFusionRules Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker =
  Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
tryOptimizeSOAC Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker
    TryFusion FusedSOAC -> TryFusion FusedSOAC -> TryFusion FusedSOAC
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
tryOptimizeKernel Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker
    TryFusion FusedSOAC -> TryFusion FusedSOAC -> TryFusion FusedSOAC
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
fuseSOACwithKer Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker
    TryFusion FusedSOAC -> TryFusion FusedSOAC -> TryFusion FusedSOAC
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
tryExposeInputs Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker

-- | Whether we are doing horizontal or vertical fusion.  Note that
-- vertical also includes "diagonal" fusion, where some producer
-- results are also produced by the final SOAC.
data Mode = Horizontal | Vertical

-- | Attempt fusing the producer into the consumer.
attemptFusion ::
  (HasScope SOACS m, MonadFreshNames m) =>
  Mode ->
  -- | Outputs of the producer that should still be output by the
  -- fusion result (corresponding to "diagonal fusion").
  Names ->
  -- | The outputs of the SOAC.
  [VName] ->
  SOAC ->
  FusedSOAC ->
  m (Maybe FusedSOAC)
attemptFusion :: forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> m (Maybe FusedSOAC)
attemptFusion Mode
mode Names
unfus_nms [VName]
outVars SOAC SOACS
soac FusedSOAC
ker = do
  scope <- m (Scope SOACS)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  tryFusion (applyFusionRules mode unfus_nms outVars soac ker) scope

-- | Check that the consumer uses all the outputs of the producer unmodified.
mapWriteFusionOK :: [VName] -> FusedSOAC -> Bool
mapWriteFusionOK :: [VName] -> FusedSOAC -> Bool
mapWriteFusionOK [VName]
outVars FusedSOAC
ker = (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
inpIds) [VName]
outVars
  where
    inpIds :: [VName]
inpIds = (Input -> Maybe VName) -> [Input] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Input -> Maybe VName
SOAC.isVarishInput (FusedSOAC -> [Input]
inputs FusedSOAC
ker)

-- | The brain of this module: Fusing a SOAC with a Kernel.
fuseSOACwithKer ::
  Mode ->
  Names ->
  [VName] ->
  SOAC ->
  FusedSOAC ->
  TryFusion FusedSOAC
fuseSOACwithKer :: Mode
-> Names
-> [VName]
-> SOAC SOACS
-> FusedSOAC
-> TryFusion FusedSOAC
fuseSOACwithKer Mode
mode Names
unfus_set [VName]
outVars SOAC SOACS
soac_p0 FusedSOAC
ker = do
  -- soac_p0 is not removed from the program until much later, so the copy of it
  -- that ends up inside the fused SOAC must be given fresh names.
  soac_p <- (Lambda SOACS -> SOAC SOACS -> SOAC SOACS)
-> SOAC SOACS -> Lambda SOACS -> SOAC SOACS
forall a b c. (a -> b -> c) -> b -> a -> c
flip Lambda SOACS -> SOAC SOACS -> SOAC SOACS
forall rep. Lambda rep -> SOAC rep -> SOAC rep
SOAC.setLambda SOAC SOACS
soac_p0 (Lambda SOACS -> SOAC SOACS)
-> TryFusion (Lambda SOACS) -> TryFusion (SOAC SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> TryFusion (Lambda SOACS)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda (SOAC SOACS -> Lambda SOACS
forall rep. SOAC rep -> Lambda rep
SOAC.lambda SOAC SOACS
soac_p0)
  -- We are fusing soac_p into soac_c, i.e, the output of soac_p is going
  -- into soac_c.
  let soac_c = FusedSOAC -> SOAC SOACS
fsSOAC FusedSOAC
ker
      inp_p_arr = SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
SOAC.inputs SOAC SOACS
soac_p
      inp_c_arr = SOAC SOACS -> [Input]
forall rep. SOAC rep -> [Input]
SOAC.inputs SOAC SOACS
soac_c
      out_p = [VName]
outVars
      out_c = FusedSOAC -> [VName]
fsOutNames FusedSOAC
ker
      lam_p = SOAC SOACS -> Lambda SOACS
forall rep. SOAC rep -> Lambda rep
SOAC.lambda SOAC SOACS
soac_p
      lam_c = SOAC SOACS -> Lambda SOACS
forall rep. SOAC rep -> Lambda rep
SOAC.lambda SOAC SOACS
soac_c
      w = SOAC SOACS -> SubExp
forall rep. SOAC rep -> SubExp
SOAC.width SOAC SOACS
soac_p
      returned_outvars = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Names -> Bool
`nameIn` Names
unfus_set) [VName]
outVars
      success [VName]
res_outnms SOAC SOACS
res_soac =
        FusedSOAC -> TryFusion FusedSOAC
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FusedSOAC -> TryFusion FusedSOAC)
-> FusedSOAC -> TryFusion FusedSOAC
forall a b. (a -> b) -> a -> b
$ FusedSOAC
ker {fsSOAC = res_soac, fsOutNames = res_outnms}

  -- Can only fuse SOACs with same width.
  guard $ SOAC.width soac_p == SOAC.width soac_c

  -- If we are getting rid of a producer output, then it must be used
  -- exclusively without any transformations.
  let ker_inputs = (Input -> VName) -> [Input] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Input -> VName
SOAC.inputArray (FusedSOAC -> [Input]
inputs FusedSOAC
ker)
      okInput VName
v Input
inp = VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= Input -> VName
SOAC.inputArray Input
inp Bool -> Bool -> Bool
|| Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Input -> Maybe VName
SOAC.isVarishInput Input
inp)
      inputOrUnfus VName
v = (Input -> Bool) -> [Input] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Input -> Bool
okInput VName
v) (FusedSOAC -> [Input]
inputs FusedSOAC
ker) Bool -> Bool -> Bool
|| VName
v VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
ker_inputs

  guard $ all inputOrUnfus outVars

  outPairs <- forM (zip outVars $ map rowType $ SOAC.typeOf soac_p) $ \(VName
outVar, Type
t) -> do
    outVar' <- Name -> TryFusion VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> TryFusion VName) -> Name -> TryFusion VName
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName VName
outVar Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_elem"
    pure (outVar, Ident outVar' t)

  let mapLikeFusionCheck =
        let (Lambda SOACS
res_lam, [Input]
new_inp) = Names
-> Lambda SOACS
-> [Input]
-> [(VName, Ident)]
-> Lambda SOACS
-> [Input]
-> (Lambda SOACS, [Input])
forall rep.
Buildable rep =>
Names
-> Lambda rep
-> [Input]
-> [(VName, Ident)]
-> Lambda rep
-> [Input]
-> (Lambda rep, [Input])
fuseMaps Names
unfus_set Lambda SOACS
lam_p [Input]
inp_p_arr [(VName, Ident)]
outPairs Lambda SOACS
lam_c [Input]
inp_c_arr
            ([VName]
extra_nms, [Type]
extra_rtps) =
              [(VName, Type)] -> ([VName], [Type])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, Type)] -> ([VName], [Type]))
-> [(VName, Type)] -> ([VName], [Type])
forall a b. (a -> b) -> a -> b
$
                ((VName, Type) -> Bool) -> [(VName, Type)] -> [(VName, Type)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Names -> Bool
`nameIn` Names
unfus_set) (VName -> Bool)
-> ((VName, Type) -> VName) -> (VName, Type) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Type) -> VName
forall a b. (a, b) -> a
fst) ([(VName, Type)] -> [(VName, Type)])
-> [(VName, Type)] -> [(VName, Type)]
forall a b. (a -> b) -> a -> b
$
                  [VName] -> [Type] -> [(VName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
outVars ([Type] -> [(VName, Type)]) -> [Type] -> [(VName, Type)]
forall a b. (a -> b) -> a -> b
$
                    (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Type -> Type
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray Int
1) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$
                      SOAC SOACS -> [Type]
forall rep. SOAC rep -> [Type]
SOAC.typeOf SOAC SOACS
soac_p
            res_lam' :: Lambda SOACS
res_lam' = Lambda SOACS
res_lam {lambdaReturnType = lambdaReturnType res_lam ++ extra_rtps}
         in ([VName]
extra_nms, Lambda SOACS
res_lam', [Input]
new_inp)

  case (soac_c, soac_p, mode) of
    (SOAC SOACS, SOAC SOACS, Mode)
_ | SOAC SOACS -> SubExp
forall rep. SOAC rep -> SubExp
SOAC.width SOAC SOACS
soac_p SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
/= SOAC SOACS -> SubExp
forall rep. SOAC rep -> SubExp
SOAC.width SOAC SOACS
soac_c -> String -> TryFusion FusedSOAC
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"SOAC widths must match."
    (SOAC SOACS
_, SOAC SOACS
_, Mode
Horizontal)
      | Bool -> Bool
not (ArrayTransforms -> Bool
SOAC.nullTransforms (ArrayTransforms -> Bool) -> ArrayTransforms -> Bool
forall a b. (a -> b) -> a -> b
$ FusedSOAC -> ArrayTransforms
fsOutputTransform FusedSOAC
ker) ->
          String -> TryFusion FusedSOAC
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Horizontal fusion is invalid in the presence of output transforms."
    (SOAC SOACS
_, SOAC SOACS
_, Mode
Vertical)
      | Names
unfus_set Names -> Names -> Bool
forall a. Eq a => a -> a -> Bool
/= Names
forall a. Monoid a => a
mempty,
        Bool -> Bool
not (ArrayTransforms -> Bool
SOAC.nullTransforms (ArrayTransforms -> Bool) -> ArrayTransforms -> Bool
forall a b. (a -> b) -> a -> b
$ FusedSOAC -> ArrayTransforms
fsOutputTransform FusedSOAC
ker) ->
          String -> TryFusion FusedSOAC
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail
            String
"Cannot perform diagonal fusion in the presence of output transforms."
    ( SOAC.Screma SubExp
_ [Input]
inp_c ScremaForm SOACS
form_c,
      SOAC.Screma SubExp
_ [Input]
inp_p ScremaForm SOACS
form_p,
      Mode
_
      ) ->
        do
          (inp, form, out) <- SubExp
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> TryFusion ([Input], ScremaForm SOACS, [VName])
forall (m :: * -> *).
(MonadFail m, MonadFreshNames m, HasScope SOACS m) =>
SubExp
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> m ([Input], ScremaForm SOACS, [VName])
fuseScrema SubExp
w [Input]
inp_p ScremaForm SOACS
form_p [VName]
out_p [Input]
inp_c ScremaForm SOACS
form_c [VName]
out_c
          success out $ SOAC.Screma w inp form
          TryFusion FusedSOAC -> TryFusion FusedSOAC -> TryFusion FusedSOAC
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> do
            -- If these two scremas cannot be fused, then see what happens if we
            -- turn the producer into a stream. The most common (only?) case of
            -- this mattering is when the producer is a scan and the consumer is
            -- a reduction.
            Just _ <- Maybe ([Scan SOACS], Lambda SOACS)
-> TryFusion (Maybe ([Scan SOACS], Lambda SOACS))
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe ([Scan SOACS], Lambda SOACS)
 -> TryFusion (Maybe ([Scan SOACS], Lambda SOACS)))
-> Maybe ([Scan SOACS], Lambda SOACS)
-> TryFusion (Maybe ([Scan SOACS], Lambda SOACS))
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> Maybe ([Scan SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
Futhark.isScanomapSOAC ScremaForm SOACS
form_p
            (soac_p', newacc_ids) <- SOAC.soacToStream soac_p
            if soac_p' /= soac_p
              then
                fuseSOACwithKer
                  mode
                  (namesFromList (map identName newacc_ids) <> unfus_set)
                  (map identName newacc_ids ++ outVars)
                  soac_p'
                  ker
              else fail "SOAC could not be turned into stream."

    -- Map-Hist fusion.
    --
    -- The 'inplace' mechanism for kernels already takes care of
    -- checking that the Hist is not writing to any array used in
    -- the Map.
    ( SOAC.Hist SubExp
_ [Input]
_ [HistOp SOACS]
ops Lambda SOACS
_,
      SOAC.Screma SubExp
_ [Input]
_ ScremaForm SOACS
form,
      Mode
_
      )
        | Maybe (Lambda SOACS) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Lambda SOACS) -> Bool) -> Maybe (Lambda SOACS) -> Bool
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form,
          -- 1. all arrays produced by the map are ONLY used (consumed)
          --    by the hist, i.e., not used elsewhere.
          (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Names -> Bool
`notNameIn` Names
unfus_set) [VName]
outVars,
          -- 2. all arrays produced by the map are input to the scatter.
          [VName] -> FusedSOAC -> Bool
mapWriteFusionOK [VName]
outVars FusedSOAC
ker -> do
            let ([VName]
extra_nms, Lambda SOACS
res_lam', [Input]
new_inp) = ([VName], Lambda SOACS, [Input])
mapLikeFusionCheck
            [VName] -> SOAC SOACS -> TryFusion FusedSOAC
success (FusedSOAC -> [VName]
fsOutNames FusedSOAC
ker [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
extra_nms) (SOAC SOACS -> TryFusion FusedSOAC)
-> SOAC SOACS -> TryFusion FusedSOAC
forall a b. (a -> b) -> a -> b
$
              SubExp -> [Input] -> [HistOp SOACS] -> Lambda SOACS -> SOAC SOACS
forall rep.
SubExp -> [Input] -> [HistOp rep] -> Lambda rep -> SOAC rep
SOAC.Hist SubExp
w [Input]
new_inp [HistOp SOACS]
ops Lambda SOACS
res_lam'

    -- Hist-Hist fusion
    ( SOAC.Hist SubExp
_ [Input]
_ [HistOp SOACS]
ops_c Lambda SOACS
_,
      SOAC.Hist SubExp
_ [Input]
_ [HistOp SOACS]
ops_p Lambda SOACS
_,
      Mode
Horizontal
      ) -> do
        let p_num_is :: Int
p_num_is = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (HistOp SOACS -> Int) -> [HistOp SOACS] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> (HistOp SOACS -> Shape) -> HistOp SOACS -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp SOACS -> Shape
forall rep. HistOp rep -> Shape
histShape) [HistOp SOACS]
ops_p
            c_num_is :: Int
c_num_is = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (HistOp SOACS -> Int) -> [HistOp SOACS] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> (HistOp SOACS -> Shape) -> HistOp SOACS -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp SOACS -> Shape
forall rep. HistOp rep -> Shape
histShape) [HistOp SOACS]
ops_c
            (Body SOACS
body_p, Body SOACS
body_c) = (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam_p, Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam_c)
            body' :: Body SOACS
body' =
              Body
                { bodyDec :: BodyDec SOACS
bodyDec = Body SOACS -> BodyDec SOACS
forall rep res. GBody rep res -> BodyDec rep
bodyDec Body SOACS
body_p, -- body_p and body_c have the same decorations
                  bodyStms :: Stms SOACS
bodyStms = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body_p Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body_c,
                  bodyResult :: [SubExpRes]
bodyResult =
                    Int -> [SubExpRes] -> [SubExpRes]
forall a. Int -> [a] -> [a]
take Int
c_num_is (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
body_c)
                      [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. [a] -> [a] -> [a]
++ Int -> [SubExpRes] -> [SubExpRes]
forall a. Int -> [a] -> [a]
take Int
p_num_is (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
body_p)
                      [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. [a] -> [a] -> [a]
++ Int -> [SubExpRes] -> [SubExpRes]
forall a. Int -> [a] -> [a]
drop Int
c_num_is (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
body_c)
                      [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. [a] -> [a] -> [a]
++ Int -> [SubExpRes] -> [SubExpRes]
forall a. Int -> [a] -> [a]
drop Int
p_num_is (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
body_p)
                }
            lam' :: Lambda SOACS
lam' =
              Lambda
                { lambdaParams :: [LParam SOACS]
lambdaParams = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam_c [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam_p,
                  lambdaBody :: Body SOACS
lambdaBody = Body SOACS
body',
                  lambdaReturnType :: [Type]
lambdaReturnType =
                    Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate (Int
c_num_is Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p_num_is) (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64)
                      [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop Int
c_num_is (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam_c)
                      [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop Int
p_num_is (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam_p)
                }
        [VName] -> SOAC SOACS -> TryFusion FusedSOAC
success (FusedSOAC -> [VName]
fsOutNames FusedSOAC
ker [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
returned_outvars) (SOAC SOACS -> TryFusion FusedSOAC)
-> SOAC SOACS -> TryFusion FusedSOAC
forall a b. (a -> b) -> a -> b
$
          SubExp -> [Input] -> [HistOp SOACS] -> Lambda SOACS -> SOAC SOACS
forall rep.
SubExp -> [Input] -> [HistOp rep] -> Lambda rep -> SOAC rep
SOAC.Hist SubExp
w ([Input]
inp_c_arr [Input] -> [Input] -> [Input]
forall a. Semigroup a => a -> a -> a
<> [Input]
inp_p_arr) ([HistOp SOACS]
ops_c [HistOp SOACS] -> [HistOp SOACS] -> [HistOp SOACS]
forall a. Semigroup a => a -> a -> a
<> [HistOp SOACS]
ops_p) Lambda SOACS
lam'
    (SOAC SOACS
_, SOAC.Hist {}, Mode
_) ->
      String -> TryFusion FusedSOAC
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Cannot fuse a Hist with anything else than a Hist or a Map"
    (SOAC.Hist {}, SOAC SOACS
_, Mode
_) ->
      String -> TryFusion FusedSOAC
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Cannot fuse a Hist with anything else than a Hist or a Map"
    ----------------------------
    -- Stream-Stream Fusions: --
    ----------------------------
    (SOAC.Stream {}, SOAC.Stream {}, Mode
_) -> do
      -- fuse two SEQUENTIAL streams
      (res_nms, res_stream) <- [VName]
-> Names
-> [VName]
-> [(VName, Ident)]
-> SOAC SOACS
-> SOAC SOACS
-> TryFusion ([VName], SOAC SOACS)
fuseStreamHelper (FusedSOAC -> [VName]
fsOutNames FusedSOAC
ker) Names
unfus_set [VName]
outVars [(VName, Ident)]
outPairs SOAC SOACS
soac_c SOAC SOACS
soac_p
      success res_nms res_stream
    -------------------------------------------------------------------
    --- If one is a stream, translate the other to a stream as well.---
    --- This does not get in trouble (infinite computation) because ---
    ---   scan's translation to Stream introduces a hindrance to    ---
    ---   (horizontal fusion), hence repeated application is for the---
    ---   moment impossible. However, if with a dependence-graph rep---
    ---   we could run in an infinite recursion, i.e., repeatedly   ---
    ---   fusing map o scan into an infinity of Stream levels!      ---
    -------------------------------------------------------------------
    (SOAC.Stream {}, SOAC SOACS
_, Mode
_) -> do
      -- If this rule is matched then soac_p is NOT a stream.
      -- To fuse a stream kernel, we transform soac_p to a stream, which
      -- borrows the sequential/parallel property of the soac_c Stream,
      -- and recursively perform stream-stream fusion.
      (soac_p', newacc_ids) <- SOAC SOACS -> TryFusion (SOAC SOACS, [Ident])
forall rep (m :: * -> *).
(HasScope rep m, MonadFreshNames m, Buildable rep, BuilderOps rep,
 Op rep ~ SOAC rep) =>
SOAC rep -> m (SOAC rep, [Ident])
SOAC.soacToStream SOAC SOACS
soac_p
      fuseSOACwithKer
        mode
        (namesFromList (map identName newacc_ids) <> unfus_set)
        (map identName newacc_ids ++ outVars)
        soac_p'
        ker
    (SOAC SOACS
_, SOAC.Stream {}, Mode
_) -> do
      -- If it reached this case then soac_c is NOT a Stream kernel,
      -- hence transform the kernel's soac to a stream and attempt
      -- stream-stream fusion recursivelly.
      -- The newly created stream corresponding to soac_c borrows the
      -- sequential/parallel property of the soac_p stream.
      (soac_c', newacc_ids) <- SOAC SOACS -> TryFusion (SOAC SOACS, [Ident])
forall rep (m :: * -> *).
(HasScope rep m, MonadFreshNames m, Buildable rep, BuilderOps rep,
 Op rep ~ SOAC rep) =>
SOAC rep -> m (SOAC rep, [Ident])
SOAC.soacToStream SOAC SOACS
soac_c
      if soac_c' /= soac_c
        then
          fuseSOACwithKer
            mode
            (namesFromList (map identName newacc_ids) <> unfus_set)
            outVars
            soac_p
            $ ker {fsSOAC = soac_c', fsOutNames = map identName newacc_ids ++ fsOutNames ker}
        else fail "SOAC could not be turned into stream."

fuseStreamHelper ::
  [VName] ->
  Names ->
  [VName] ->
  [(VName, Ident)] ->
  SOAC ->
  SOAC ->
  TryFusion ([VName], SOAC)
fuseStreamHelper :: [VName]
-> Names
-> [VName]
-> [(VName, Ident)]
-> SOAC SOACS
-> SOAC SOACS
-> TryFusion ([VName], SOAC SOACS)
fuseStreamHelper
  [VName]
out_kernms
  Names
unfus_set
  [VName]
outVars
  [(VName, Ident)]
outPairs
  (SOAC.Stream SubExp
w2 [Input]
inp2_arr [SubExp]
nes2 Lambda SOACS
lam2)
  (SOAC.Stream SubExp
_ [Input]
inp1_arr [SubExp]
nes1 Lambda SOACS
lam1) = do
    -- very similar to redomap o redomap composition, but need
    -- to remove first the `chunk' parameters of streams'
    -- lambdas and put them in the resulting stream lambda.
    let chunk1 :: Param Type
chunk1 = [Param Type] -> Param Type
forall a. HasCallStack => [a] -> a
head ([Param Type] -> Param Type) -> [Param Type] -> Param Type
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam1
        chunk2 :: Param Type
chunk2 = [Param Type] -> Param Type
forall a. HasCallStack => [a] -> a
head ([Param Type] -> Param Type) -> [Param Type] -> Param Type
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam2
        hmnms :: Map VName VName
hmnms = [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
chunk2, Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
chunk1)]
        lam20 :: Lambda SOACS
lam20 = Map VName VName -> Lambda SOACS -> Lambda SOACS
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
hmnms Lambda SOACS
lam2
        lam1' :: Lambda SOACS
lam1' = Lambda SOACS
lam1 {lambdaParams = tail $ lambdaParams lam1}
        lam2' :: Lambda SOACS
lam2' = Lambda SOACS
lam20 {lambdaParams = tail $ lambdaParams lam20}
        (Lambda SOACS
res_lam', [Input]
new_inp) =
          Names
-> [VName]
-> Lambda SOACS
-> [SubExp]
-> [SubExp]
-> [Input]
-> [(VName, Ident)]
-> Lambda SOACS
-> [SubExp]
-> [SubExp]
-> [Input]
-> (Lambda SOACS, [Input])
forall rep.
Buildable rep =>
Names
-> [VName]
-> Lambda rep
-> [SubExp]
-> [SubExp]
-> [Input]
-> [(VName, Ident)]
-> Lambda rep
-> [SubExp]
-> [SubExp]
-> [Input]
-> (Lambda rep, [Input])
fuseRedomap
            Names
unfus_set
            [VName]
outVars
            Lambda SOACS
lam1'
            []
            [SubExp]
nes1
            [Input]
inp1_arr
            [(VName, Ident)]
outPairs
            Lambda SOACS
lam2'
            []
            [SubExp]
nes2
            [Input]
inp2_arr
        res_lam'' :: Lambda SOACS
res_lam'' = Lambda SOACS
res_lam' {lambdaParams = chunk1 : lambdaParams res_lam'}
        unfus_accs :: [VName]
unfus_accs = Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
take ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes1) [VName]
outVars
        unfus_arrs :: [VName]
unfus_arrs = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
unfus_accs) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Names -> Bool
`nameIn` Names
unfus_set) [VName]
outVars
    ([VName], SOAC SOACS) -> TryFusion ([VName], SOAC SOACS)
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( [VName]
unfus_accs [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
out_kernms [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
unfus_arrs,
        SubExp -> [Input] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. SubExp -> [Input] -> [SubExp] -> Lambda rep -> SOAC rep
SOAC.Stream SubExp
w2 [Input]
new_inp ([SubExp]
nes1 [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
nes2) Lambda SOACS
res_lam''
      )
fuseStreamHelper [VName]
_ Names
_ [VName]
_ [(VName, Ident)]
_ SOAC SOACS
_ SOAC SOACS
_ = String -> TryFusion ([VName], SOAC SOACS)
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Cannot Fuse Streams!"

-- Here follows optimizations and transforms to expose fusability.

optimizeKernel :: Maybe [VName] -> FusedSOAC -> TryFusion FusedSOAC
optimizeKernel :: Maybe [VName] -> FusedSOAC -> TryFusion FusedSOAC
optimizeKernel Maybe [VName]
inp FusedSOAC
ker = do
  (soac, resTrans) <- Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
optimizeSOAC Maybe [VName]
inp (FusedSOAC -> SOAC SOACS
fsSOAC FusedSOAC
ker) (FusedSOAC -> ArrayTransforms
fsOutputTransform FusedSOAC
ker)
  pure $ ker {fsSOAC = soac, fsOutputTransform = resTrans}

optimizeSOAC ::
  Maybe [VName] ->
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)
optimizeSOAC :: Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
optimizeSOAC Maybe [VName]
inp SOAC SOACS
soac ArrayTransforms
os = do
  res <- ((Bool, SOAC SOACS, ArrayTransforms)
 -> (Maybe [VName]
     -> SOAC SOACS
     -> ArrayTransforms
     -> TryFusion (SOAC SOACS, ArrayTransforms))
 -> TryFusion (Bool, SOAC SOACS, ArrayTransforms))
-> (Bool, SOAC SOACS, ArrayTransforms)
-> [Maybe [VName]
    -> SOAC SOACS
    -> ArrayTransforms
    -> TryFusion (SOAC SOACS, ArrayTransforms)]
-> TryFusion (Bool, SOAC SOACS, ArrayTransforms)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Bool, SOAC SOACS, ArrayTransforms)
-> (Maybe [VName]
    -> SOAC SOACS
    -> ArrayTransforms
    -> TryFusion (SOAC SOACS, ArrayTransforms))
-> TryFusion (Bool, SOAC SOACS, ArrayTransforms)
comb (Bool
False, SOAC SOACS
soac, ArrayTransforms
os) [Maybe [VName]
 -> SOAC SOACS
 -> ArrayTransforms
 -> TryFusion (SOAC SOACS, ArrayTransforms)]
optimizations
  case res of
    (Bool
False, SOAC SOACS
_, ArrayTransforms
_) -> String -> TryFusion (SOAC SOACS, ArrayTransforms)
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"No optimisation applied"
    (Bool
True, SOAC SOACS
soac', ArrayTransforms
os') -> (SOAC SOACS, ArrayTransforms)
-> TryFusion (SOAC SOACS, ArrayTransforms)
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SOAC SOACS
soac', ArrayTransforms
os')
  where
    comb :: (Bool, SOAC SOACS, ArrayTransforms)
-> (Maybe [VName]
    -> SOAC SOACS
    -> ArrayTransforms
    -> TryFusion (SOAC SOACS, ArrayTransforms))
-> TryFusion (Bool, SOAC SOACS, ArrayTransforms)
comb (Bool
changed, SOAC SOACS
soac', ArrayTransforms
os') Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
f =
      do
        (soac'', os'') <- Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
f Maybe [VName]
inp SOAC SOACS
soac' ArrayTransforms
os
        pure (True, soac'', os'')
        TryFusion (Bool, SOAC SOACS, ArrayTransforms)
-> TryFusion (Bool, SOAC SOACS, ArrayTransforms)
-> TryFusion (Bool, SOAC SOACS, ArrayTransforms)
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (Bool, SOAC SOACS, ArrayTransforms)
-> TryFusion (Bool, SOAC SOACS, ArrayTransforms)
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
changed, SOAC SOACS
soac', ArrayTransforms
os')

type Optimization =
  Maybe [VName] ->
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)

optimizations :: [Optimization]
optimizations :: [Maybe [VName]
 -> SOAC SOACS
 -> ArrayTransforms
 -> TryFusion (SOAC SOACS, ArrayTransforms)]
optimizations = [Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
iswim, Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
unflattenAccOnlyMap]

iswim ::
  Maybe [VName] ->
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)
iswim :: Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
iswim Maybe [VName]
_ (SOAC.Screma SubExp
w [Input]
arrs ScremaForm SOACS
form) ArrayTransforms
ots
  | Just [Futhark.Scan Lambda SOACS
scan_fun [SubExp]
nes] <- ScremaForm SOACS -> Maybe [Scan SOACS]
forall rep. ScremaForm rep -> Maybe [Scan rep]
Futhark.isScanSOAC ScremaForm SOACS
form,
    Just (Pat Type
map_pat, StmAux ()
map_aux, SubExp
map_w, Lambda SOACS
map_fun) <- Lambda SOACS -> Maybe (Pat Type, StmAux (), SubExp, Lambda SOACS)
rwimPossible Lambda SOACS
scan_fun,
    Just [VName]
nes_names <- (SubExp -> Maybe VName) -> [SubExp] -> Maybe [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> Maybe VName
subExpVar [SubExp]
nes = do
      let nes_idents :: [Ident]
nes_idents = (VName -> Type -> Ident) -> [VName] -> [Type] -> [Ident]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Ident
Ident [VName]
nes_names ([Type] -> [Ident]) -> [Type] -> [Ident]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
scan_fun
          map_nes :: [Input]
map_nes = (Ident -> Input) -> [Ident] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> Input
SOAC.identInput [Ident]
nes_idents
          map_arrs' :: [Input]
map_arrs' = [Input]
map_nes [Input] -> [Input] -> [Input]
forall a. [a] -> [a] -> [a]
++ (Input -> Input) -> [Input] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Input -> Input
SOAC.transposeInput Int
0 Int
1) [Input]
arrs
          ([Param Type]
scan_acc_params, [Param Type]
scan_elem_params) =
            Int -> [Param Type] -> ([Param Type], [Param Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Input] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Input]
arrs) ([Param Type] -> ([Param Type], [Param Type]))
-> [Param Type] -> ([Param Type], [Param Type])
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
scan_fun
          map_params :: [Param Type]
map_params =
            (Param Type -> Param Type) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Param Type
LParam SOACS -> LParam SOACS
removeParamOuterDim [Param Type]
scan_acc_params
              [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ (Param Type -> Param Type) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> LParam SOACS -> LParam SOACS
setParamOuterDimTo SubExp
w) [Param Type]
scan_elem_params
          map_rettype :: [Type]
map_rettype = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> SubExp -> Type
forall d m.
TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
`setOuterSize` SubExp
w) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
scan_fun

          scan_params :: [LParam SOACS]
scan_params = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
map_fun
          scan_body :: Body SOACS
scan_body = Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
map_fun
          scan_rettype :: [Type]
scan_rettype = Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
map_fun
          scan_fun' :: Lambda SOACS
scan_fun' = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [LParam SOACS]
scan_params [Type]
scan_rettype Body SOACS
scan_body
          nes' :: [SubExp]
nes' = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var ([VName] -> [SubExp]) -> [VName] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
take ([Input] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Input]
map_nes) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
map_params
          arrs' :: [VName]
arrs' = Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
drop ([Input] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Input]
map_nes) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
map_params

      scan_form <- [Scan SOACS] -> TryFusion (ScremaForm SOACS)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
scanSOAC [Lambda SOACS -> [SubExp] -> Scan SOACS
forall rep. Lambda rep -> [SubExp] -> Scan rep
Futhark.Scan Lambda SOACS
scan_fun' [SubExp]
nes']

      let map_body =
            Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody
              ( Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm (Stm SOACS -> Stms SOACS) -> Stm SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$
                  Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let (SubExp -> Pat Type -> Pat Type
setPatOuterDimTo SubExp
w Pat Type
map_pat) (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp SOACS -> Stm SOACS)
-> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> Stm SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op SOACS -> Exp SOACS
SOAC SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (SOAC SOACS -> Stm SOACS) -> SOAC SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
                    SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Futhark.Screma SubExp
w [VName]
arrs' ScremaForm SOACS
scan_form
              )
              ([SubExpRes] -> Body SOACS) -> [SubExpRes] -> Body SOACS
forall a b. (a -> b) -> a -> b
$ [VName] -> [SubExpRes]
varsRes
              ([VName] -> [SubExpRes]) -> [VName] -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
map_pat
          map_fun' = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
map_params [Type]
map_rettype Body SOACS
map_body
          perm = case Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
scan_fun of -- instead of map_fun
            [] -> []
            Type
t : [Type]
_ -> Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int
2 .. Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
t]

      (,ots SOAC.|> SOAC.Rearrange map_aux perm)
        . SOAC.Screma map_w map_arrs'
        <$> mapSOAC map_fun'
iswim Maybe [VName]
_ SOAC SOACS
_ ArrayTransforms
_ =
  String -> TryFusion (SOAC SOACS, ArrayTransforms)
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"ISWIM does not apply."

-- | When a pure-map Screma returns exclusively accumulator results and some
-- non-accumulator inputs carry a 2D-to-1D flattening Reshape transform, we
-- can "unflatten" the map:
--
--   Screma(n*m, {flat_a1:[n*m]t1, ..., acc_p:acc(...)}, lam)
--   where lam : (t1, ..., acc) → acc
--
-- becomes
--
--   Screma(n, {a1:[n][m]t1, ..., acc_p:acc(...)}, outer_lam)
--   where outer_lam = \(row_a1:[m]t1, ..., acc_p:acc(...)) →
--     Screma(m, {row_a1, ..., acc_p}, lam)
--
-- This exposes the 2D inputs directly, enabling the standard fusion rules to
-- fuse the resulting Screma(n,...) with an upstream Screma(n,...) producer.
unflattenAccOnlyMap ::
  Maybe [VName] ->
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)
unflattenAccOnlyMap :: Maybe [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
unflattenAccOnlyMap (Just [VName]
outVars) (SOAC.Screma SubExp
_nm [Input]
inps ScremaForm SOACS
form) ArrayTransforms
ots = do
  lam <- Maybe (Lambda SOACS) -> TryFusion (Lambda SOACS)
forall a. Maybe a -> TryFusion a
liftMaybe (Maybe (Lambda SOACS) -> TryFusion (Lambda SOACS))
-> Maybe (Lambda SOACS) -> TryFusion (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form
  -- All results must be accumulator types.
  guard $ all isAcc $ lambdaReturnType lam
  -- Only apply when the producer outputs non-scalar rows (rank > 1), meaning
  -- pullReshape cannot handle this case (it requires scalar-leaf map nests).
  -- When the producer outputs scalars, the simpler prepend approach works fine.
  outVarTypes <- mapM lookupType outVars
  guard $ any ((> 1) . arrayRank) outVarTypes
  -- Partition inputs paired with their lambda params: those with a 2D→1D
  -- flattening Reshape vs. those that pass through unchanged (acc params).
  -- A flattening reshape: base type is 2D, first transform collapses it to 1D.
  let classifyInp (inp :: Input
inp@(SOAC.Input ArrayTransforms
ts VName
_v Type
base_t), b
p)
        | SOAC.Reshape StmAux ()
_aux NewShape SubExp
ns SOAC.:< ArrayTransforms
ts' <- ArrayTransforms -> ViewF
SOAC.viewf ArrayTransforms
ts,
          Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
base_t Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2,
          Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (NewShape SubExp -> Shape
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
ns) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 =
            (Input, b) -> Either (Input, b) (Input, b)
forall a b. a -> Either a b
Left (ArrayTransforms -> VName -> Type -> Input
SOAC.Input ArrayTransforms
ts' VName
_v Type
base_t, b
p)
        | Bool
otherwise =
            (Input, b) -> Either (Input, b) (Input, b)
forall a b. b -> Either a b
Right (Input
inp, b
p)
      (flat_pairs, pass_pairs) =
        partitionEithers $ zipWith (curry classifyInp) inps (lambdaParams lam)
  -- Need at least one flattened input.
  guard $ not (null flat_pairs)
  -- The non-flattened inputs must be accumulators, because we are changing the
  -- width of the SOAC.
  guard $ all (isAcc . SOAC.inputType . fst) pass_pairs
  -- All flattened inputs must agree on the outer dim n and inner dim m.
  let dims2d TypeBase Shape o
base_t = (Int -> TypeBase Shape o -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 TypeBase Shape o
base_t, Int -> TypeBase Shape o -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
1 TypeBase Shape o
base_t)
      getBaseTy (SOAC.Input ArrayTransforms
_ VName
_ Type
base_t, b
_) = Type
base_t
      (n, m) = dims2d (getBaseTy (head flat_pairs))
  guard $ all ((== (n, m)) . dims2d . getBaseTy) flat_pairs
  -- The lambda params for the flat inputs get their type changed from [n*m]t
  -- to [m]t (a single row).  Pass-through params are unchanged.
  let mkRowParam (a
_, Param (TypeBase (ShapeBase d) m)
p) = Param (TypeBase (ShapeBase d) m)
p {paramDec = rowType (paramDec p)}
      flat_row_params = ((Input, Param Type) -> Param Type)
-> [(Input, Param Type)] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (Input, Param Type) -> Param Type
forall {a} {d} {m}.
(a, Param (TypeBase (ShapeBase d) m))
-> Param (TypeBase (ShapeBase d) m)
mkRowParam [(Input, Param Type)]
flat_pairs
      pass_params = ((Input, Param Type) -> Param Type)
-> [(Input, Param Type)] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (Input, Param Type) -> Param Type
forall a b. (a, b) -> b
snd [(Input, Param Type)]
pass_pairs
      inner_lam_params = [Param Type]
flat_row_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
pass_params
  inner_lam <- renameLambda $ lam {lambdaParams = inner_lam_params}
  inner_form <- mapSOAC inner_lam
  -- Inner Screma over m: plain-variable inputs for the row params, then
  -- plain-variable inputs for the pass-through (acc) params.
  let inner_inps =
        (Param Type -> Input) -> [Param Type] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (Ident -> Input
SOAC.identInput (Ident -> Input) -> (Param Type -> Ident) -> Param Type -> Input
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> Ident
paramToIdent) [Param Type]
flat_row_params
          [Input] -> [Input] -> [Input]
forall a. [a] -> [a] -> [a]
++ (Param Type -> Input) -> [Param Type] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (Ident -> Input
SOAC.identInput (Ident -> Input) -> (Param Type -> Ident) -> Param Type -> Input
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> Ident
paramToIdent) [Param Type]
pass_params
      inner_soac = SubExp -> [Input] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [Input] -> ScremaForm rep -> SOAC rep
SOAC.Screma SubExp
m [Input]
inner_inps ScremaForm SOACS
inner_form
  -- Outer lambda: same param names but outer params have type [m]t (rows).
  let outer_lam_params = [Param Type]
flat_row_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
pass_params
  outer_lam <- runLambdaBuilder outer_lam_params $ do
    inner_exp <- SOAC.toExp inner_soac
    res <- letTupExp "inner_acc" inner_exp
    pure $ map (subExpRes . Var) res
  outer_form <- mapSOAC outer_lam
  -- Outer Screma over n: 2D inputs (flatten reshape stripped) then pass-through.
  let outer_inps = ((Input, Param Type) -> Input) -> [(Input, Param Type)] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (Input, Param Type) -> Input
forall a b. (a, b) -> a
fst [(Input, Param Type)]
flat_pairs [Input] -> [Input] -> [Input]
forall a. [a] -> [a] -> [a]
++ ((Input, Param Type) -> Input) -> [(Input, Param Type)] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (Input, Param Type) -> Input
forall a b. (a, b) -> a
fst [(Input, Param Type)]
pass_pairs
  pure (SOAC.Screma n outer_inps outer_form, ots)
unflattenAccOnlyMap Maybe [VName]
_ SOAC SOACS
_ ArrayTransforms
_ =
  String -> TryFusion (SOAC SOACS, ArrayTransforms)
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unflattenAccOnlyMap does not apply."

paramToIdent :: Param Type -> Ident
paramToIdent :: Param Type -> Ident
paramToIdent Param Type
p = VName -> Type -> Ident
Ident (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p) (Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p)

removeParamOuterDim :: LParam SOACS -> LParam SOACS
removeParamOuterDim :: LParam SOACS -> LParam SOACS
removeParamOuterDim LParam SOACS
param =
  let t :: Type
t = Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
LParam SOACS
param
   in LParam SOACS
param {paramDec = t}

setParamOuterDimTo :: SubExp -> LParam SOACS -> LParam SOACS
setParamOuterDimTo :: SubExp -> LParam SOACS -> LParam SOACS
setParamOuterDimTo SubExp
w LParam SOACS
param =
  let t :: Type
t = Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
LParam SOACS
param Type -> SubExp -> Type
forall d m.
TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
`setOuterSize` SubExp
w
   in LParam SOACS
param {paramDec = t}

setPatOuterDimTo :: SubExp -> Pat Type -> Pat Type
setPatOuterDimTo :: SubExp -> Pat Type -> Pat Type
setPatOuterDimTo SubExp
w = (Type -> Type) -> Pat Type -> Pat Type
forall a b. (a -> b) -> Pat a -> Pat b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type -> SubExp -> Type
forall d m.
TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
`setOuterSize` SubExp
w)

-- Now for fiddling with transpositions...

commonTransforms ::
  [VName] ->
  [SOAC.Input] ->
  (SOAC.ArrayTransforms, [SOAC.Input])
commonTransforms :: [VName] -> [Input] -> (ArrayTransforms, [Input])
commonTransforms [VName]
interesting [Input]
inps = [(Bool, Input)] -> (ArrayTransforms, [Input])
commonTransforms' [(Bool, Input)]
inps'
  where
    inps' :: [(Bool, Input)]
inps' =
      [ (Input -> VName
SOAC.inputArray Input
inp VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
interesting, Input
inp)
      | Input
inp <- [Input]
inps
      ]

commonTransforms' :: [(Bool, SOAC.Input)] -> (SOAC.ArrayTransforms, [SOAC.Input])
commonTransforms' :: [(Bool, Input)] -> (ArrayTransforms, [Input])
commonTransforms' [(Bool, Input)]
inps =
  case ((Maybe ArrayTransform, [(Bool, Input)])
 -> (Bool, Input) -> Maybe (Maybe ArrayTransform, [(Bool, Input)]))
-> (Maybe ArrayTransform, [(Bool, Input)])
-> [(Bool, Input)]
-> Maybe (Maybe ArrayTransform, [(Bool, Input)])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Maybe ArrayTransform, [(Bool, Input)])
-> (Bool, Input) -> Maybe (Maybe ArrayTransform, [(Bool, Input)])
inspect (Maybe ArrayTransform
forall a. Maybe a
Nothing, []) [(Bool, Input)]
inps of
    Just (Just ArrayTransform
mot, [(Bool, Input)]
inps') -> (ArrayTransforms -> ArrayTransforms)
-> (ArrayTransforms, [Input]) -> (ArrayTransforms, [Input])
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (ArrayTransform
mot ArrayTransform -> ArrayTransforms -> ArrayTransforms
SOAC.<|) ((ArrayTransforms, [Input]) -> (ArrayTransforms, [Input]))
-> (ArrayTransforms, [Input]) -> (ArrayTransforms, [Input])
forall a b. (a -> b) -> a -> b
$ [(Bool, Input)] -> (ArrayTransforms, [Input])
commonTransforms' ([(Bool, Input)] -> (ArrayTransforms, [Input]))
-> [(Bool, Input)] -> (ArrayTransforms, [Input])
forall a b. (a -> b) -> a -> b
$ [(Bool, Input)] -> [(Bool, Input)]
forall a. [a] -> [a]
reverse [(Bool, Input)]
inps'
    Maybe (Maybe ArrayTransform, [(Bool, Input)])
_ -> (ArrayTransforms
SOAC.noTransforms, ((Bool, Input) -> Input) -> [(Bool, Input)] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (Bool, Input) -> Input
forall a b. (a, b) -> b
snd [(Bool, Input)]
inps)
  where
    -- Two reshapes with the same shape are compatible even if their certs differ;
    -- merge the certs to produce a single common reshape transform.
    compatibleTransforms :: ArrayTransform -> ArrayTransform -> Maybe ArrayTransform
compatibleTransforms (SOAC.Reshape StmAux ()
aux1 NewShape SubExp
shape1) (SOAC.Reshape StmAux ()
aux2 NewShape SubExp
shape2)
      | NewShape SubExp
shape1 NewShape SubExp -> NewShape SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== NewShape SubExp
shape2 =
          ArrayTransform -> Maybe ArrayTransform
forall a. a -> Maybe a
Just (ArrayTransform -> Maybe ArrayTransform)
-> ArrayTransform -> Maybe ArrayTransform
forall a b. (a -> b) -> a -> b
$ StmAux () -> NewShape SubExp -> ArrayTransform
SOAC.Reshape (StmAux ()
aux1 StmAux () -> StmAux () -> StmAux ()
forall a. Semigroup a => a -> a -> a
<> StmAux ()
aux2) NewShape SubExp
shape1
    compatibleTransforms ArrayTransform
ot1 ArrayTransform
ot2
      | ArrayTransform
ot1 ArrayTransform -> ArrayTransform -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayTransform
ot2 = ArrayTransform -> Maybe ArrayTransform
forall a. a -> Maybe a
Just ArrayTransform
ot1
    compatibleTransforms ArrayTransform
_ ArrayTransform
_ = Maybe ArrayTransform
forall a. Maybe a
Nothing

    inspect :: (Maybe ArrayTransform, [(Bool, Input)])
-> (Bool, Input) -> Maybe (Maybe ArrayTransform, [(Bool, Input)])
inspect (Maybe ArrayTransform
mot, [(Bool, Input)]
prev) (Bool
True, Input
inp) =
      case (Maybe ArrayTransform
mot, Input -> Maybe (ArrayTransform, Input)
inputToOutput Input
inp) of
        (Maybe ArrayTransform
Nothing, Just (ArrayTransform
ot, Input
inp')) -> (Maybe ArrayTransform, [(Bool, Input)])
-> Maybe (Maybe ArrayTransform, [(Bool, Input)])
forall a. a -> Maybe a
Just (ArrayTransform -> Maybe ArrayTransform
forall a. a -> Maybe a
Just ArrayTransform
ot, (Bool
True, Input
inp') (Bool, Input) -> [(Bool, Input)] -> [(Bool, Input)]
forall a. a -> [a] -> [a]
: [(Bool, Input)]
prev)
        (Just ArrayTransform
ot1, Just (ArrayTransform
ot2, Input
inp'))
          | Just ArrayTransform
combined <- ArrayTransform -> ArrayTransform -> Maybe ArrayTransform
compatibleTransforms ArrayTransform
ot1 ArrayTransform
ot2 ->
              (Maybe ArrayTransform, [(Bool, Input)])
-> Maybe (Maybe ArrayTransform, [(Bool, Input)])
forall a. a -> Maybe a
Just (ArrayTransform -> Maybe ArrayTransform
forall a. a -> Maybe a
Just ArrayTransform
combined, (Bool
True, Input
inp') (Bool, Input) -> [(Bool, Input)] -> [(Bool, Input)]
forall a. a -> [a] -> [a]
: [(Bool, Input)]
prev)
        (Maybe ArrayTransform, Maybe (ArrayTransform, Input))
_ -> Maybe (Maybe ArrayTransform, [(Bool, Input)])
forall a. Maybe a
Nothing
    inspect (Maybe ArrayTransform
mot, [(Bool, Input)]
prev) (Bool, Input)
inp = (Maybe ArrayTransform, [(Bool, Input)])
-> Maybe (Maybe ArrayTransform, [(Bool, Input)])
forall a. a -> Maybe a
Just (Maybe ArrayTransform
mot, (Bool, Input)
inp (Bool, Input) -> [(Bool, Input)] -> [(Bool, Input)]
forall a. a -> [a] -> [a]
: [(Bool, Input)]
prev)

mapDepth :: MapNest -> Int
mapDepth :: MapNest -> Int
mapDepth (MapNest.MapNest SubExp
_ Lambda SOACS
lam [Nesting]
levels [Input]
_) =
  Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
resDims ([Nesting] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Nesting]
levels) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  where
    resDims :: Int
resDims = [Type] -> Int
forall {shape} {o}. ArrayShape shape => [TypeBase shape o] -> Int
minDim ([Type] -> Int) -> [Type] -> Int
forall a b. (a -> b) -> a -> b
$ case [Nesting]
levels of
      [] -> Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam
      Nesting
nest : [Nesting]
_ -> Nesting -> [Type]
MapNest.nestingReturnType Nesting
nest
    minDim :: [TypeBase shape o] -> Int
minDim [] = Int
0
    minDim (TypeBase shape o
t : [TypeBase shape o]
ts) = (Int -> Int -> Int) -> Int -> [Int] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (TypeBase shape o -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank TypeBase shape o
t) ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (TypeBase shape o -> Int) -> [TypeBase shape o] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase shape o -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank [TypeBase shape o]
ts

pullRearrange ::
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)
pullRearrange :: SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullRearrange SOAC SOACS
soac ArrayTransforms
ots = do
  nest <- Maybe MapNest -> TryFusion MapNest
forall a. Maybe a -> TryFusion a
liftMaybe (Maybe MapNest -> TryFusion MapNest)
-> TryFusion (Maybe MapNest) -> TryFusion MapNest
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SOAC SOACS -> TryFusion (Maybe MapNest)
forall (m :: * -> *).
(MonadFreshNames m, LocalScope SOACS m) =>
SOAC SOACS -> m (Maybe MapNest)
MapNest.fromSOAC SOAC SOACS
soac
  SOAC.Rearrange cs perm SOAC.:< ots' <- pure $ SOAC.viewf ots
  if rearrangeReach perm <= mapDepth nest
    then do
      let -- Expand perm to cover the full extent of the input dimensionality
          perm' Input
inp = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
r [Int]
perm [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
perm .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            where
              r :: Int
r = Input -> Int
SOAC.inputRank Input
inp
          addPerm Input
inp = ArrayTransform -> Input -> Input
SOAC.addTransform (StmAux () -> [Int] -> ArrayTransform
SOAC.Rearrange StmAux ()
cs ([Int] -> ArrayTransform) -> [Int] -> ArrayTransform
forall a b. (a -> b) -> a -> b
$ Input -> [Int]
perm' Input
inp) Input
inp
          inputs' = (Input -> Input) -> [Input] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map Input -> Input
addPerm ([Input] -> [Input]) -> [Input] -> [Input]
forall a b. (a -> b) -> a -> b
$ MapNest -> [Input]
MapNest.inputs MapNest
nest
      soac' <-
        MapNest.toSOAC $
          inputs' `MapNest.setInputs` rearrangeReturnTypes nest perm
      pure (soac', ots')
    else fail "Cannot pull transpose"

pullIndex ::
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)
pullIndex :: SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullIndex (SOAC.Screma SubExp
_ [Input]
inps ScremaForm SOACS
form) ArrayTransforms
ots
  | SOAC.Index StmAux ()
cs slice :: Slice SubExp
slice@(Slice (ds :: DimIndex SubExp
ds@(DimSlice SubExp
_ SubExp
w' SubExp
_) : [DimIndex SubExp]
inner_ds))
      SOAC.:< ArrayTransforms
ots' <-
      ArrayTransforms -> ViewF
SOAC.viewf ArrayTransforms
ots,
    Just Lambda SOACS
lam <- ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form = do
      let sliceInput :: Input -> Input
sliceInput Input
inp =
            ArrayTransform -> Input -> Input
SOAC.addTransform
              (StmAux () -> Slice SubExp -> ArrayTransform
SOAC.Index StmAux ()
cs (Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice (Input -> Type
SOAC.inputType Input
inp) [DimIndex SubExp
ds]))
              Input
inp
          sliceRes :: SubExpRes -> BuilderT SOACS (State VNameSource) SubExpRes
sliceRes (SubExpRes Certs
rcs (Var VName
v)) =
            Certs
-> BuilderT SOACS (State VNameSource) SubExpRes
-> BuilderT SOACS (State VNameSource) SubExpRes
forall a.
Certs
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
rcs
              (BuilderT SOACS (State VNameSource) SubExpRes
 -> BuilderT SOACS (State VNameSource) SubExpRes)
-> (Exp SOACS -> BuilderT SOACS (State VNameSource) SubExpRes)
-> Exp SOACS
-> BuilderT SOACS (State VNameSource) SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> SubExpRes)
-> BuilderT SOACS (State VNameSource) SubExp
-> BuilderT SOACS (State VNameSource) SubExpRes
forall a b.
(a -> b)
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> SubExpRes
subExpRes
              (BuilderT SOACS (State VNameSource) SubExp
 -> BuilderT SOACS (State VNameSource) SubExpRes)
-> (Exp SOACS -> BuilderT SOACS (State VNameSource) SubExp)
-> Exp SOACS
-> BuilderT SOACS (State VNameSource) SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_sliced")
              (Exp SOACS -> BuilderT SOACS (State VNameSource) SubExpRes)
-> Exp SOACS -> BuilderT SOACS (State VNameSource) SubExpRes
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (VName -> Slice SubExp -> BasicOp
Index VName
v ([DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice [DimIndex SubExp]
inner_ds))
          sliceRes SubExpRes
r = SubExpRes -> BuilderT SOACS (State VNameSource) SubExpRes
forall a. a -> BuilderT SOACS (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExpRes
r
          inner_changed :: Bool
inner_changed =
            (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
              ((Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> Shape -> Shape
forall d. Int -> ShapeBase d -> ShapeBase d
stripDims Int
1 (Slice SubExp -> Shape
forall d. Slice d -> ShapeBase d
sliceShape Slice SubExp
slice)) (Shape -> Bool) -> (Type -> Shape) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape)
              (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam)
      lam' <-
        if Bool -> Bool
not Bool
inner_changed
          then Lambda SOACS -> TryFusion (Lambda SOACS)
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda SOACS
lam
          else
            [LParam SOACS]
-> Builder SOACS [SubExpRes] -> TryFusion (Lambda SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
[LParam rep] -> Builder rep [SubExpRes] -> m (Lambda rep)
runLambdaBuilder (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam) (Builder SOACS [SubExpRes] -> TryFusion (Lambda SOACS))
-> Builder SOACS [SubExpRes] -> TryFusion (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$
              (SubExpRes -> BuilderT SOACS (State VNameSource) SubExpRes)
-> [SubExpRes] -> Builder SOACS [SubExpRes]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExpRes -> BuilderT SOACS (State VNameSource) SubExpRes
sliceRes ([SubExpRes] -> Builder SOACS [SubExpRes])
-> Builder SOACS [SubExpRes] -> Builder SOACS [SubExpRes]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Rep (BuilderT SOACS (State VNameSource)))
-> Builder SOACS [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam)
      (,ots') . SOAC.Screma w' (map sliceInput inps) <$> mapSOAC lam'
pullIndex SOAC SOACS
_ ArrayTransforms
_ = String -> TryFusion (SOAC SOACS, ArrayTransforms)
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Cannot pull index"

pushRearrange ::
  [VName] ->
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)
pushRearrange :: [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
pushRearrange [VName]
inpIds SOAC SOACS
soac ArrayTransforms
ots = do
  nest <- Maybe MapNest -> TryFusion MapNest
forall a. Maybe a -> TryFusion a
liftMaybe (Maybe MapNest -> TryFusion MapNest)
-> TryFusion (Maybe MapNest) -> TryFusion MapNest
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SOAC SOACS -> TryFusion (Maybe MapNest)
forall (m :: * -> *).
(MonadFreshNames m, LocalScope SOACS m) =>
SOAC SOACS -> m (Maybe MapNest)
MapNest.fromSOAC SOAC SOACS
soac
  (perm, inputs') <- liftMaybe $ fixupInputs inpIds $ MapNest.inputs nest
  if rearrangeReach perm <= mapDepth nest
    then do
      let invertRearrange = StmAux () -> [Int] -> ArrayTransform
SOAC.Rearrange StmAux ()
forall a. Monoid a => a
mempty ([Int] -> ArrayTransform) -> [Int] -> ArrayTransform
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int]
rearrangeInverse [Int]
perm
      soac' <-
        MapNest.toSOAC $
          inputs'
            `MapNest.setInputs` rearrangeReturnTypes nest perm
      pure (soac', invertRearrange SOAC.<| ots)
    else fail "Cannot push transpose"

-- | Actually also rearranges indices.
rearrangeReturnTypes :: MapNest -> [Int] -> MapNest
rearrangeReturnTypes :: MapNest -> [Int] -> MapNest
rearrangeReturnTypes nest :: MapNest
nest@(MapNest.MapNest SubExp
w Lambda SOACS
body [Nesting]
nestings [Input]
inps) [Int]
perm =
  SubExp -> Lambda SOACS -> [Nesting] -> [Input] -> MapNest
MapNest.MapNest
    SubExp
w
    Lambda SOACS
body
    ( (Nesting -> [Type] -> Nesting)
-> [Nesting] -> [[Type]] -> [Nesting]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
        Nesting -> [Type] -> Nesting
setReturnType
        [Nesting]
nestings
        ([[Type]] -> [Nesting]) -> [[Type]] -> [Nesting]
forall a b. (a -> b) -> a -> b
$ Int -> [[Type]] -> [[Type]]
forall a. Int -> [a] -> [a]
drop Int
1
        ([[Type]] -> [[Type]]) -> [[Type]] -> [[Type]]
forall a b. (a -> b) -> a -> b
$ ([Type] -> [Type]) -> [Type] -> [[Type]]
forall a. (a -> a) -> a -> [a]
iterate ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType) [Type]
ts
    )
    [Input]
inps
  where
    origts :: [Type]
origts = MapNest -> [Type]
MapNest.typeOf MapNest
nest
    -- The permutation may be deeper than the rank of the type,
    -- but it is required that it is an identity permutation
    -- beyond that.  This is supposed to be checked as an
    -- invariant by whoever calls rearrangeReturnTypes.
    rearrangeType' :: Type -> Type
rearrangeType' Type
t = [Int] -> Type -> Type
forall d m.
[Int] -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rearrangeType (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
t) [Int]
perm) Type
t
    ts :: [Type]
ts = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
rearrangeType' [Type]
origts

    setReturnType :: Nesting -> [Type] -> Nesting
setReturnType Nesting
nesting [Type]
t' =
      Nesting
nesting {MapNest.nestingReturnType = t'}

fixupInputs :: [VName] -> [SOAC.Input] -> Maybe ([Int], [SOAC.Input])
fixupInputs :: [VName] -> [Input] -> Maybe ([Int], [Input])
fixupInputs [VName]
inpIds [Input]
inps =
  case (Input -> Maybe [Int]) -> [Input] -> [[Int]]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Input -> Maybe [Int]
inputRearrange ([Input] -> [[Int]]) -> [Input] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ (Input -> Bool) -> [Input] -> [Input]
forall a. (a -> Bool) -> [a] -> [a]
filter Input -> Bool
exposable [Input]
inps of
    [Int]
perm : [[Int]]
_ -> do
      inps' <- (Input -> Maybe Input) -> [Input] -> Maybe [Input]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> [Int] -> Input -> Maybe Input
fixupInput ([Int] -> Int
rearrangeReach [Int]
perm) [Int]
perm) [Input]
inps
      pure (perm, inps')
    [[Int]]
_ -> Maybe ([Int], [Input])
forall a. Maybe a
Nothing
  where
    exposable :: Input -> Bool
exposable = (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
inpIds) (VName -> Bool) -> (Input -> VName) -> Input -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> VName
SOAC.inputArray

    inputRearrange :: Input -> Maybe [Int]
inputRearrange (SOAC.Input ArrayTransforms
ts VName
_ Type
_)
      | ArrayTransforms
_ SOAC.:> SOAC.Rearrange StmAux ()
_ [Int]
perm <- ArrayTransforms -> ViewL
SOAC.viewl ArrayTransforms
ts = [Int] -> Maybe [Int]
forall a. a -> Maybe a
Just [Int]
perm
    inputRearrange Input
_ = Maybe [Int]
forall a. Maybe a
Nothing

    fixupInput :: Int -> [Int] -> Input -> Maybe Input
fixupInput Int
d [Int]
perm Input
inp
      | Int
r <- Input -> Int
SOAC.inputRank Input
inp,
        Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
d =
          Input -> Maybe Input
forall a. a -> Maybe a
Just (Input -> Maybe Input) -> Input -> Maybe Input
forall a b. (a -> b) -> a -> b
$ ArrayTransform -> Input -> Input
SOAC.addTransform (StmAux () -> [Int] -> ArrayTransform
SOAC.Rearrange StmAux ()
forall a. Monoid a => a
mempty ([Int] -> ArrayTransform) -> [Int] -> ArrayTransform
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
r [Int]
perm) Input
inp
      | Bool
otherwise = Maybe Input
forall a. Maybe a
Nothing

pullReshape :: SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms)
pullReshape :: SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullReshape SOAC SOACS
soac ArrayTransforms
ots = do
  Just mapnest <- SOAC SOACS -> TryFusion (Maybe MapNest)
forall (m :: * -> *).
(MonadFreshNames m, LocalScope SOACS m) =>
SOAC SOACS -> m (Maybe MapNest)
MapNest.fromSOAC SOAC SOACS
soac
  SOAC.Reshape cs newshape SOAC.:< ots' <- pure $ SOAC.viewf ots
  -- This handles only the easy case where the underlying lambda is
  -- scalar. The more complicated cases could also be handled, but
  -- requires more tricky checks.
  guard $
    all
      ((== MapNest.depth mapnest) . arrayRank)
      (MapNest.typeOf mapnest)
  -- Reshaping the nest changes its parallel dimensions, e.g. an unflattening
  -- splits one dimension into several. When the innermost lambda contains a
  -- SOAC (such as a reduction), this can inhibit later optimisations: a matrix
  -- multiplication whose result is unflattened becomes a nest where one operand
  -- varies with two of the dimensions, which tiling cannot handle (see
  -- tests/tiling/tiling_mm_unflatten.fut). In that case we only allow
  -- coercions. As a simple heuristic, we only look at the top-level statements
  -- of the innermost lambda. The proper solution is for tiling to handle such
  -- nests.
  guard $
    reshapeKind newshape == ReshapeCoerce
      || not (any (isSOAC . stmExp) (bodyStms (lambdaBody (MapNest.mapNestLambda mapnest))))
  mapnest' <- MapNest.reshape cs newshape mapnest
  soac' <- MapNest.toSOAC mapnest'
  pure (soac', ots')
  where
    isSOAC :: Exp rep -> Bool
isSOAC Op {} = Bool
True
    isSOAC Exp rep
_ = Bool
False

-- Tie it all together in exposeInputs (for making inputs to a
-- consumer available) and pullOutputTransforms (for moving
-- output-transforms of a producer to its inputs instead).

exposeInputs ::
  [VName] ->
  FusedSOAC ->
  TryFusion (FusedSOAC, SOAC.ArrayTransforms)
exposeInputs :: [VName] -> FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs [VName]
inpIds FusedSOAC
ker =
  (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs' (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms))
-> TryFusion FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TryFusion FusedSOAC
pushRearrange')
    TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs' (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms))
-> TryFusion FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TryFusion FusedSOAC
pullRearrange')
    TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs' (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms))
-> TryFusion FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TryFusion FusedSOAC
pullReshape')
    TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs' (FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms))
-> TryFusion FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TryFusion FusedSOAC
pullIndex')
    TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs' FusedSOAC
ker
  where
    ot :: ArrayTransforms
ot = FusedSOAC -> ArrayTransforms
fsOutputTransform FusedSOAC
ker

    pushRearrange' :: TryFusion FusedSOAC
pushRearrange' = do
      (soac', ot') <- [VName]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
pushRearrange [VName]
inpIds (FusedSOAC -> SOAC SOACS
fsSOAC FusedSOAC
ker) ArrayTransforms
ot
      pure
        ker
          { fsSOAC = soac',
            fsOutputTransform = ot'
          }

    pullRearrange' :: TryFusion FusedSOAC
pullRearrange' = do
      (soac', ot') <- SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullRearrange (FusedSOAC -> SOAC SOACS
fsSOAC FusedSOAC
ker) ArrayTransforms
ot
      unless (SOAC.nullTransforms ot') $
        fail "pullRearrange was not enough"
      pure
        ker
          { fsSOAC = soac',
            fsOutputTransform = SOAC.noTransforms
          }

    pullReshape' :: TryFusion FusedSOAC
pullReshape' = do
      (soac', ot') <- SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullReshape (FusedSOAC -> SOAC SOACS
fsSOAC FusedSOAC
ker) ArrayTransforms
ot
      unless (SOAC.nullTransforms ot') $
        fail "pullReshape was not enough"
      pure
        ker
          { fsSOAC = soac',
            fsOutputTransform = SOAC.noTransforms
          }

    pullIndex' :: TryFusion FusedSOAC
pullIndex' = do
      (soac', ot') <- SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullIndex (FusedSOAC -> SOAC SOACS
fsSOAC FusedSOAC
ker) ArrayTransforms
ot
      unless (SOAC.nullTransforms ot') $
        fail "pullIndex was not enough"
      pure
        ker
          { fsSOAC = soac',
            fsOutputTransform = SOAC.noTransforms
          }

    exposeInputs' :: FusedSOAC -> TryFusion (FusedSOAC, ArrayTransforms)
exposeInputs' FusedSOAC
ker' =
      case [VName] -> [Input] -> (ArrayTransforms, [Input])
commonTransforms [VName]
inpIds ([Input] -> (ArrayTransforms, [Input]))
-> [Input] -> (ArrayTransforms, [Input])
forall a b. (a -> b) -> a -> b
$ FusedSOAC -> [Input]
inputs FusedSOAC
ker' of
        (ArrayTransforms
ot', [Input]
inps')
          | (Input -> Bool) -> [Input] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Input -> Bool
exposed [Input]
inps' ->
              (FusedSOAC, ArrayTransforms)
-> TryFusion (FusedSOAC, ArrayTransforms)
forall a. a -> TryFusion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FusedSOAC
ker' {fsSOAC = inps' `SOAC.setInputs` fsSOAC ker'}, ArrayTransforms
ot')
        (ArrayTransforms, [Input])
_ -> String -> TryFusion (FusedSOAC, ArrayTransforms)
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Cannot expose"

    exposed :: Input -> Bool
exposed (SOAC.Input ArrayTransforms
ts VName
_ Type
_)
      | ArrayTransforms -> Bool
SOAC.nullTransforms ArrayTransforms
ts = Bool
True
    exposed Input
inp = Input -> VName
SOAC.inputArray Input
inp VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
inpIds

outputTransformPullers ::
  [ SOAC ->
    SOAC.ArrayTransforms ->
    TryFusion (SOAC, SOAC.ArrayTransforms)
  ]
outputTransformPullers :: [SOAC SOACS
 -> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)]
outputTransformPullers = [SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullRearrange, SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullReshape, SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullIndex]

pullOutputTransforms ::
  SOAC ->
  SOAC.ArrayTransforms ->
  TryFusion (SOAC, SOAC.ArrayTransforms)
pullOutputTransforms :: SOAC SOACS
-> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)
pullOutputTransforms = [SOAC SOACS
 -> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)]
-> SOAC SOACS
-> ArrayTransforms
-> TryFusion (SOAC SOACS, ArrayTransforms)
forall {t} {t}.
[t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)]
-> t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)
attempt [SOAC SOACS
 -> ArrayTransforms -> TryFusion (SOAC SOACS, ArrayTransforms)]
outputTransformPullers
  where
    attempt :: [t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)]
-> t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)
attempt [] t
_ t
_ = String -> TryFusion (SOAC SOACS, ArrayTransforms)
forall a. String -> TryFusion a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Cannot pull anything"
    attempt (t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)
p : [t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)]
ps) t
soac t
ots =
      do
        (soac', ots') <- t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)
p t
soac t
ots
        if SOAC.nullTransforms ots'
          then pure (soac', SOAC.noTransforms)
          else pullOutputTransforms soac' ots' <|> pure (soac', ots')
        TryFusion (SOAC SOACS, ArrayTransforms)
-> TryFusion (SOAC SOACS, ArrayTransforms)
-> TryFusion (SOAC SOACS, ArrayTransforms)
forall a. TryFusion a -> TryFusion a -> TryFusion a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> [t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)]
-> t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)
attempt [t -> t -> TryFusion (SOAC SOACS, ArrayTransforms)]
ps t
soac t
ots