{-# LANGUAGE CPP                  #-}
{-# LANGUAGE ConstraintKinds      #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE InstanceSigs         #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE PatternGuards        #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE ViewPatterns         #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.Trafo.Fusion
-- Copyright   : [2012..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- This module implements producer/producer and consumer/producer fusion as a
-- term rewriting of the Accelerate AST.
--
-- The function 'quench' perform the source-to-source fusion transformation,
-- while 'anneal' additionally makes the representation of embedded producers
-- explicit by representing the AST as a 'DelayedAcc' of manifest and delayed
-- nodes.
--

module Data.Array.Accelerate.Trafo.Fusion (

  convertAcc,  convertAccWith,
  convertAfun, convertAfunWith,

) where

import Data.BitSet
import Data.Array.Accelerate.AST
import Data.Array.Accelerate.AST.LeftHandSide
import Data.Array.Accelerate.AST.Environment
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Trafo.Config
import Data.Array.Accelerate.Trafo.Var
import Data.Array.Accelerate.Trafo.Delayed
import Data.Array.Accelerate.Trafo.Environment
import Data.Array.Accelerate.Trafo.Shrink
import Data.Array.Accelerate.Trafo.Simplify
import Data.Array.Accelerate.Trafo.Substitution
import Data.Array.Accelerate.Representation.Array                   ( Array, ArrayR(..), ArraysR )
import Data.Array.Accelerate.Representation.Shape                   ( ShapeR(..), shapeType )
import Data.Array.Accelerate.Representation.Slice
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Analysis.Match
import Data.Array.Accelerate.Type

import Data.Array.Accelerate.Debug.Internal.Flags                   ( array_fusion )
import qualified Data.Array.Accelerate.Debug.Internal.Stats         as Stats
#ifdef ACCELERATE_DEBUG
import System.IO.Unsafe -- for debugging
#endif

import Data.Function
import Lens.Micro                                                 ( over, mapped, _2 )
import Prelude                                                      hiding ( exp, until )


-- Delayed Array Fusion
-- ====================

-- | Apply the fusion transformation to a closed de Bruijn AST
--
convertAcc :: HasCallStack => Acc arrs -> DelayedAcc arrs
convertAcc :: forall arrs. HasCallStack => Acc arrs -> DelayedAcc arrs
convertAcc = Config -> Acc arrs -> DelayedAcc arrs
forall arrs. HasCallStack => Config -> Acc arrs -> DelayedAcc arrs
convertAccWith Config
defaultOptions

convertAccWith :: HasCallStack => Config -> Acc arrs -> DelayedAcc arrs
convertAccWith :: forall arrs. HasCallStack => Config -> Acc arrs -> DelayedAcc arrs
convertAccWith Config
config = DelayedAcc arrs -> DelayedAcc arrs
forall a. a -> a
withSimplStats (DelayedAcc arrs -> DelayedAcc arrs)
-> (Acc arrs -> DelayedAcc arrs) -> Acc arrs -> DelayedAcc arrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Config -> Acc arrs -> DelayedAcc arrs
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
convertOpenAcc Config
config

-- | Apply the fusion transformation to a function of array arguments
--
convertAfun :: HasCallStack => Afun f -> DelayedAfun f
convertAfun :: forall f. HasCallStack => Afun f -> DelayedAfun f
convertAfun = Config -> Afun f -> DelayedAfun f
forall f. HasCallStack => Config -> Afun f -> DelayedAfun f
convertAfunWith Config
defaultOptions

convertAfunWith :: HasCallStack => Config -> Afun f -> DelayedAfun f
convertAfunWith :: forall f. HasCallStack => Config -> Afun f -> DelayedAfun f
convertAfunWith Config
config = DelayedAfun f -> DelayedAfun f
forall a. a -> a
withSimplStats (DelayedAfun f -> DelayedAfun f)
-> (Afun f -> DelayedAfun f) -> Afun f -> DelayedAfun f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Config -> Afun f -> DelayedAfun f
forall aenv f.
HasCallStack =>
Config -> OpenAfun aenv f -> DelayedOpenAfun aenv f
convertOpenAfun Config
config

-- -- | Apply the fusion transformation to the array computations embedded
-- --   in a sequence computation.
--
-- convertSeq :: Bool -> Seq a -> DelayedSeq a
-- convertSeq fuseAcc (embedSeq (embedOpenAcc fuseAcc) -> ExtendSeq aenv s)
--   = withSimplStats (DelayedSeq (cvtE aenv) (convertOpenSeq fuseAcc s))
--   where
--     cvtE :: Extend OpenAcc aenv aenv' -> Extend DelayedOpenAcc aenv aenv'
--     cvtE BaseEnv                                          = BaseEnv
--     cvtE (PushEnv env a) | a' <- convertOpenAcc fuseAcc a = PushEnv (cvtE env) a'

withSimplStats :: a -> a
#ifdef ACCELERATE_DEBUG
withSimplStats x = unsafePerformIO Stats.resetSimplCount `seq` x
#else
withSimplStats :: forall a. a -> a
withSimplStats a
x = a
x
#endif


-- | Apply the fusion transformation to an AST. This consists of two phases:
--
--    1. A bottom-up traversal that converts nodes into the internal delayed
--       representation, merging adjacent producer/producer pairs.
--
--    2. A top-down traversal that makes the representation of fused
--       consumer/producer pairs explicit as a 'DelayedAcc' of manifest and
--       delayed nodes.
--
-- TLM: Note that there really is no ambiguity as to which state an array will
--      be in following this process: an array will be either delayed or
--      manifest, and the two helper functions are even named as such! We should
--      encode this property in the type somehow...
--
convertOpenAcc
    :: HasCallStack
    => Config
    -> OpenAcc aenv arrs
    -> DelayedOpenAcc aenv arrs
convertOpenAcc :: forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
convertOpenAcc Config
config = Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config (OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs)
-> (OpenAcc aenv arrs -> OpenAcc aenv arrs)
-> OpenAcc aenv arrs
-> DelayedOpenAcc aenv arrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Embed OpenAcc aenv arrs -> OpenAcc aenv arrs)
-> (OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> OpenAcc aenv arrs
-> OpenAcc aenv arrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
embedOpenAcc Config
config


-- Convert array computations into an embeddable delayed representation.
-- Reapply the embedding function from the first pass and unpack the
-- representation. It is safe to match on BaseEnv because the first pass
-- will put producers adjacent to the term consuming it.
--
delayed
    :: HasCallStack
    => Config
    -> OpenAcc aenv (Array sh e)
    -> DelayedOpenAcc aenv (Array sh e)
delayed :: forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config (Config
-> OpenAcc aenv (Array sh e) -> Embed OpenAcc aenv (Array sh e)
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
embedOpenAcc Config
config -> Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' (Array sh e)
cc)
  | Extend ArrayR OpenAcc aenv aenv'
BaseEnv <- Extend ArrayR OpenAcc aenv aenv'
env
  = Cunctation r aenv' (Array sh e)
-> Either
     (Cunctation M aenv' (Array sh e)) (Cunctation D aenv' (Array sh e))
forall r aenv a.
HasCallStack =>
Cunctation r aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
simplify Cunctation r aenv' (Array sh e)
cc Either
  (Cunctation M aenv' (Array sh e)) (Cunctation D aenv' (Array sh e))
-> (Either
      (Cunctation M aenv' (Array sh e)) (Cunctation D aenv' (Array sh e))
    -> DelayedOpenAcc aenv (Array sh e))
-> DelayedOpenAcc aenv (Array sh e)
forall a b. a -> (a -> b) -> b
& \case
      Left (Done ArrayVars aenv' (Array sh e)
v)                                           -> InjectAcc DelayedOpenAcc
-> ArrayVars aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
forall (acc :: * -> * -> *) aenv arrs.
InjectAcc acc -> ArrayVars aenv arrs -> acc aenv arrs
avarsIn PreOpenAcc DelayedOpenAcc env t -> DelayedOpenAcc env t
InjectAcc DelayedOpenAcc
Manifest ArrayVars aenv (Array sh e)
ArrayVars aenv' (Array sh e)
v
      Right Cunctation D aenv' (Array sh e)
d       -> Cunctation D aenv' (Array sh e)
d Cunctation D aenv' (Array sh e)
-> (Cunctation D aenv' (Array sh e)
    -> DelayedOpenAcc aenv (Array sh e))
-> DelayedOpenAcc aenv (Array sh e)
forall a b. a -> (a -> b) -> b
& \case
        Yield ArrayR (Array sh e)
aR Exp aenv' sh
sh Fun aenv' (sh -> e)
f                                         -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Fun aenv (Int -> e)
-> DelayedOpenAcc aenv (Array sh e)
forall sh e aenv.
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Fun aenv (Int -> e)
-> DelayedOpenAcc aenv (Array sh e)
Delayed ArrayR (Array sh e)
ArrayR (Array sh e)
aR Exp aenv sh
Exp aenv' sh
sh Fun aenv (sh -> e)
Fun aenv' (sh -> e)
f (OpenFun () aenv (sh -> e)
Fun aenv' (sh -> e)
f OpenFun () aenv (sh -> e)
-> OpenFun () aenv (Int -> sh) -> Fun aenv (Int -> e)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ShapeR sh -> OpenExp () aenv sh -> OpenFun () aenv (Int -> sh)
forall sh env aenv.
ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (Int -> sh)
fromIndex (ArrayR (Array sh e) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh e)
aR) OpenExp () aenv sh
Exp aenv' sh
sh)
        Step  ArrayR (Array sh' b)
aR Exp aenv' sh'
sh Fun aenv' (sh' -> sh)
p Fun aenv' (a -> b)
f ArrayVar aenv' (Array sh a)
v
          | Just sh' :~: sh
Refl <- Exp aenv' sh' -> OpenExp () aenv' sh -> Maybe (sh' :~: sh)
forall env aenv s t.
OpenExp env aenv s -> OpenExp env aenv t -> Maybe (s :~: t)
matchOpenExp Exp aenv' sh'
sh (ArrayVar aenv' (Array sh a) -> OpenExp () aenv' sh
forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape ArrayVar aenv' (Array sh a)
v)
          , Just sh' :~: sh
Refl <- Fun aenv' (sh' -> sh) -> Maybe (sh' :~: sh)
forall env aenv a b. OpenFun env aenv (a -> b) -> Maybe (a :~: b)
isIdentity Fun aenv' (sh' -> sh)
p                         -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Fun aenv (Int -> e)
-> DelayedOpenAcc aenv (Array sh e)
forall sh e aenv.
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Fun aenv (Int -> e)
-> DelayedOpenAcc aenv (Array sh e)
Delayed ArrayR (Array sh e)
ArrayR (Array sh' b)
aR Exp aenv sh
Exp aenv' sh'
sh (OpenFun () aenv (a -> e)
Fun aenv' (a -> b)
f OpenFun () aenv (a -> e)
-> OpenFun () aenv (sh -> a) -> Fun aenv (sh -> e)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ArrayVar aenv (Array sh a) -> OpenFun () aenv (sh -> a)
forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray ArrayVar aenv (Array sh a)
ArrayVar aenv' (Array sh a)
v) (OpenFun () aenv (a -> e)
Fun aenv' (a -> b)
f OpenFun () aenv (a -> e)
-> OpenFun () aenv (Int -> a) -> Fun aenv (Int -> e)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ArrayVar aenv (Array sh a) -> OpenFun () aenv (Int -> a)
forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (Int -> e)
linearIndex ArrayVar aenv (Array sh a)
ArrayVar aenv' (Array sh a)
v)
          | OpenFun () aenv' (sh' -> b)
f'        <- Fun aenv' (a -> b)
f Fun aenv' (a -> b)
-> OpenFun () aenv' (sh' -> a) -> OpenFun () aenv' (sh' -> b)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ArrayVar aenv' (Array sh a) -> Fun aenv' (sh -> a)
forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray ArrayVar aenv' (Array sh a)
v Fun aenv' (sh -> a)
-> Fun aenv' (sh' -> sh) -> OpenFun () aenv' (sh' -> a)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` Fun aenv' (sh' -> sh)
p -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Fun aenv (Int -> e)
-> DelayedOpenAcc aenv (Array sh e)
forall sh e aenv.
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Fun aenv (Int -> e)
-> DelayedOpenAcc aenv (Array sh e)
Delayed ArrayR (Array sh e)
ArrayR (Array sh' b)
aR Exp aenv sh
Exp aenv' sh'
sh Fun aenv (sh -> e)
OpenFun () aenv' (sh' -> b)
f' (OpenFun () aenv (sh' -> e)
OpenFun () aenv' (sh' -> b)
f' OpenFun () aenv (sh' -> e)
-> OpenFun () aenv (Int -> sh') -> Fun aenv (Int -> e)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ShapeR sh' -> OpenExp () aenv sh' -> OpenFun () aenv (Int -> sh')
forall sh env aenv.
ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (Int -> sh)
fromIndex (ArrayR (Array sh' b) -> ShapeR sh'
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh' b)
aR) OpenExp () aenv sh'
Exp aenv' sh'
sh)
  --
  | Bool
otherwise
  = Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config (Embed OpenAcc aenv (Array sh e) -> OpenAcc aenv (Array sh e)
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> Embed OpenAcc aenv (Array sh e)
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' (Array sh e)
cc))


-- Convert array programs as manifest terms.
--
manifest
    :: HasCallStack
    => Config
    -> OpenAcc aenv a
    -> DelayedOpenAcc aenv a
manifest :: forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config (OpenAcc PreOpenAcc OpenAcc aenv a
pacc) =
  let fusionError :: r
fusionError = Format r r -> r
forall r a. HasCallStack => Format r a -> a
internalError Format r r
"unexpected fusible materials"
  in
  PreOpenAcc DelayedOpenAcc aenv a -> DelayedOpenAcc aenv a
InjectAcc DelayedOpenAcc
Manifest (PreOpenAcc DelayedOpenAcc aenv a -> DelayedOpenAcc aenv a)
-> PreOpenAcc DelayedOpenAcc aenv a -> DelayedOpenAcc aenv a
forall a b. (a -> b) -> a -> b
$ case PreOpenAcc OpenAcc aenv a
pacc of
    -- Non-fusible terms
    -- -----------------
    Avar ArrayVar aenv (Array sh e)
ix                 -> ArrayVar aenv (Array sh e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar ArrayVar aenv (Array sh e)
ix
    Use ArrayR (Array sh e)
aR Array sh e
a                -> ArrayR (Array sh e)
-> Array sh e -> PreOpenAcc DelayedOpenAcc aenv (Array sh e)
forall sh e (acc :: * -> * -> *) aenv.
ArrayR (Array sh e)
-> Array sh e -> PreOpenAcc acc aenv (Array sh e)
Use ArrayR (Array sh e)
aR Array sh e
a
    Unit TypeR e
t Exp aenv e
e                -> TypeR e
-> Exp aenv e -> PreOpenAcc DelayedOpenAcc aenv (Array () e)
forall e aenv (acc :: * -> * -> *).
TypeR e -> Exp aenv e -> PreOpenAcc acc aenv (Array () e)
Unit TypeR e
t Exp aenv e
e
    Alet ALeftHandSide bndArrs aenv aenv'
lhs OpenAcc aenv bndArrs
bnd OpenAcc aenv' a
body       -> ALeftHandSide bndArrs aenv aenv'
-> DelayedOpenAcc aenv bndArrs
-> DelayedOpenAcc aenv' a
-> PreOpenAcc DelayedOpenAcc aenv a
forall a aenv aenv' b.
HasCallStack =>
ALeftHandSide a aenv aenv'
-> DelayedOpenAcc aenv a
-> DelayedOpenAcc aenv' b
-> PreOpenAcc DelayedOpenAcc aenv b
alet ALeftHandSide bndArrs aenv aenv'
lhs (Config -> OpenAcc aenv bndArrs -> DelayedOpenAcc aenv bndArrs
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv bndArrs
bnd) (Config -> OpenAcc aenv' a -> DelayedOpenAcc aenv' a
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv' a
body)
    Acond Exp aenv PrimBool
p OpenAcc aenv a
t OpenAcc aenv a
e             -> Exp aenv PrimBool
-> DelayedOpenAcc aenv a
-> DelayedOpenAcc aenv a
-> PreOpenAcc DelayedOpenAcc aenv a
forall aenv (acc :: * -> * -> *) a.
Exp aenv PrimBool
-> acc aenv a -> acc aenv a -> PreOpenAcc acc aenv a
Acond Exp aenv PrimBool
p (Config -> OpenAcc aenv a -> DelayedOpenAcc aenv a
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv a
t) (Config -> OpenAcc aenv a -> DelayedOpenAcc aenv a
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv a
e)
    Awhile PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
p PreOpenAfun OpenAcc aenv (a -> a)
f OpenAcc aenv a
a            -> PreOpenAfun DelayedOpenAcc aenv (a -> Scalar PrimBool)
-> PreOpenAfun DelayedOpenAcc aenv (a -> a)
-> DelayedOpenAcc aenv a
-> PreOpenAcc DelayedOpenAcc aenv a
forall (acc :: * -> * -> *) aenv a.
PreOpenAfun acc aenv (a -> Scalar PrimBool)
-> PreOpenAfun acc aenv (a -> a)
-> acc aenv a
-> PreOpenAcc acc aenv a
Awhile (PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
-> PreOpenAfun DelayedOpenAcc aenv (a -> Scalar PrimBool)
forall aenv f.
HasCallStack =>
OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
cvtAF PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
p) (PreOpenAfun OpenAcc aenv (a -> a)
-> PreOpenAfun DelayedOpenAcc aenv (a -> a)
forall aenv f.
HasCallStack =>
OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
cvtAF PreOpenAfun OpenAcc aenv (a -> a)
f) (Config -> OpenAcc aenv a -> DelayedOpenAcc aenv a
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv a
a)
    Apair OpenAcc aenv as
a1 OpenAcc aenv bs
a2             -> DelayedOpenAcc aenv as
-> DelayedOpenAcc aenv bs
-> PreOpenAcc DelayedOpenAcc aenv (as, bs)
forall (acc :: * -> * -> *) aenv as bs.
acc aenv as -> acc aenv bs -> PreOpenAcc acc aenv (as, bs)
Apair (Config -> OpenAcc aenv as -> DelayedOpenAcc aenv as
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv as
a1) (Config -> OpenAcc aenv bs -> DelayedOpenAcc aenv bs
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv bs
a2)
    PreOpenAcc OpenAcc aenv a
Anil                    -> PreOpenAcc DelayedOpenAcc aenv a
PreOpenAcc DelayedOpenAcc aenv ()
forall (acc :: * -> * -> *) aenv. PreOpenAcc acc aenv ()
Anil
    Atrace Message arrs1
msg OpenAcc aenv arrs1
a1 OpenAcc aenv a
a2        -> Message arrs1
-> DelayedOpenAcc aenv arrs1
-> DelayedOpenAcc aenv a
-> PreOpenAcc DelayedOpenAcc aenv a
forall arrs1 (acc :: * -> * -> *) aenv a.
Message arrs1
-> acc aenv arrs1 -> acc aenv a -> PreOpenAcc acc aenv a
Atrace Message arrs1
msg (Config -> OpenAcc aenv arrs1 -> DelayedOpenAcc aenv arrs1
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv arrs1
a1) (Config -> OpenAcc aenv a -> DelayedOpenAcc aenv a
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv a
a2)
    Apply ArraysR a
repr PreOpenAfun OpenAcc aenv (arrs1 -> a)
f OpenAcc aenv arrs1
a          -> ArraysR a
-> PreOpenAfun DelayedOpenAcc aenv (arrs1 -> a)
-> DelayedOpenAcc aenv arrs1
-> PreOpenAcc DelayedOpenAcc aenv a
forall b aenv a.
HasCallStack =>
ArraysR b
-> PreOpenAfun DelayedOpenAcc aenv (a -> b)
-> DelayedOpenAcc aenv a
-> PreOpenAcc DelayedOpenAcc aenv b
apply ArraysR a
repr (PreOpenAfun OpenAcc aenv (arrs1 -> a)
-> PreOpenAfun DelayedOpenAcc aenv (arrs1 -> a)
forall aenv f.
HasCallStack =>
OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
cvtAF PreOpenAfun OpenAcc aenv (arrs1 -> a)
f) (Config -> OpenAcc aenv arrs1 -> DelayedOpenAcc aenv arrs1
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv arrs1
a)
    Aforeign ArraysR a
repr asm (as -> a)
ff PreAfun OpenAcc (as -> a)
f OpenAcc aenv as
a    -> ArraysR a
-> asm (as -> a)
-> PreAfun DelayedOpenAcc (as -> a)
-> DelayedOpenAcc aenv as
-> PreOpenAcc DelayedOpenAcc aenv a
forall (asm :: * -> *) a as (acc :: * -> * -> *) aenv.
Foreign asm =>
ArraysR a
-> asm (as -> a)
-> PreAfun acc (as -> a)
-> acc aenv as
-> PreOpenAcc acc aenv a
Aforeign ArraysR a
repr asm (as -> a)
ff (PreAfun OpenAcc (as -> a) -> PreAfun DelayedOpenAcc (as -> a)
forall aenv f.
HasCallStack =>
OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
cvtAF PreAfun OpenAcc (as -> a)
f) (Config -> OpenAcc aenv as -> DelayedOpenAcc aenv as
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv as
a)

    -- Producers
    -- ---------
    --
    -- Some producers might still exist as a manifest array. Typically this
    -- is because they are the last stage of the computation, or the result
    -- of a let-binding to be used multiple times. The input array here
    -- should be a evaluated array term, else something went wrong.
    --
    Map TypeR e'
t Fun aenv (e -> e')
f OpenAcc aenv (Array sh e)
a               -> TypeR e'
-> Fun aenv (e -> e')
-> DelayedOpenAcc aenv (Array sh e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e')
forall e' aenv e (acc :: * -> * -> *) sh.
TypeR e'
-> Fun aenv (e -> e')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Map TypeR e'
t Fun aenv (e -> e')
f (Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array sh e)
a)
    Generate ArrayR (Array sh e)
repr Exp aenv sh
sh Fun aenv (sh -> e)
f      -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e)
forall sh e aenv (acc :: * -> * -> *).
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> PreOpenAcc acc aenv (Array sh e)
Generate ArrayR (Array sh e)
repr Exp aenv sh
sh Fun aenv (sh -> e)
f
    Transform ArrayR (Array sh' b)
repr Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a1 -> b)
f OpenAcc aenv (Array sh a1)
a -> ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> DelayedOpenAcc aenv (Array sh a1)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh' b)
forall sh' b aenv sh a1 (acc :: * -> * -> *).
ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> acc aenv (Array sh a1)
-> PreOpenAcc acc aenv (Array sh' b)
Transform ArrayR (Array sh' b)
repr Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a1 -> b)
f (Config
-> OpenAcc aenv (Array sh a1) -> DelayedOpenAcc aenv (Array sh a1)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array sh a1)
a)
    Backpermute ShapeR sh'
shR Exp aenv sh'
sh Fun aenv (sh' -> sh)
p OpenAcc aenv (Array sh e)
a  -> ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> DelayedOpenAcc aenv (Array sh e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh' e)
forall sh' aenv sh (acc :: * -> * -> *) e.
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Backpermute ShapeR sh'
shR Exp aenv sh'
sh Fun aenv (sh' -> sh)
p (Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array sh e)
a)
    Reshape ShapeR sh
slr Exp aenv sh
sl OpenAcc aenv (Array sh' e)
a        -> ShapeR sh
-> Exp aenv sh
-> DelayedOpenAcc aenv (Array sh' e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e)
forall sh aenv (acc :: * -> * -> *) sh' e.
ShapeR sh
-> Exp aenv sh
-> acc aenv (Array sh' e)
-> PreOpenAcc acc aenv (Array sh e)
Reshape ShapeR sh
slr Exp aenv sh
sl (Config
-> OpenAcc aenv (Array sh' e) -> DelayedOpenAcc aenv (Array sh' e)
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv (Array sh' e)
a)

    Replicate{}             -> PreOpenAcc DelayedOpenAcc aenv a
forall {r}. r
fusionError
    Slice{}                 -> PreOpenAcc DelayedOpenAcc aenv a
forall {r}. r
fusionError
    ZipWith{}               -> PreOpenAcc DelayedOpenAcc aenv a
forall {r}. r
fusionError

    -- Consumers
    -- ---------
    --
    -- Embed producers directly into the representation. For delayed terms
    -- with local bindings, these will have been floated up above the
    -- consumer already
    --
    Fold Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a              -> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> DelayedOpenAcc aenv (Array (sh, Int) e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e)
forall aenv e (acc :: * -> * -> *) sh.
Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array sh e)
Fold     Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z (Config
-> OpenAcc aenv (Array (sh, Int) e)
-> DelayedOpenAcc aenv (Array (sh, Int) e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array (sh, Int) e)
a)
    FoldSeg IntegralType i
i Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a OpenAcc aenv (Segments i)
s       -> IntegralType i
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> DelayedOpenAcc aenv (Array (sh, Int) e)
-> DelayedOpenAcc aenv (Segments i)
-> PreOpenAcc DelayedOpenAcc aenv (Array (sh, Int) e)
forall i aenv e (acc :: * -> * -> *) sh.
IntegralType i
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> acc aenv (Segments i)
-> PreOpenAcc acc aenv (Array (sh, Int) e)
FoldSeg  IntegralType i
i Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z (Config
-> OpenAcc aenv (Array (sh, Int) e)
-> DelayedOpenAcc aenv (Array (sh, Int) e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array (sh, Int) e)
a) (Config
-> OpenAcc aenv (Segments i) -> DelayedOpenAcc aenv (Segments i)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Segments i)
s)
    Scan  Direction
d Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a           -> Direction
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> DelayedOpenAcc aenv (Array (sh, Int) e)
-> PreOpenAcc DelayedOpenAcc aenv (Array (sh, Int) e)
forall aenv e (acc :: * -> * -> *) sh.
Direction
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array (sh, Int) e)
Scan     Direction
d Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z (Config
-> OpenAcc aenv (Array (sh, Int) e)
-> DelayedOpenAcc aenv (Array (sh, Int) e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array (sh, Int) e)
a)
    Scan' Direction
d Fun aenv (e -> e -> e)
f Exp aenv e
z OpenAcc aenv (Array (sh, Int) e)
a           -> Direction
-> Fun aenv (e -> e -> e)
-> Exp aenv e
-> DelayedOpenAcc aenv (Array (sh, Int) e)
-> PreOpenAcc DelayedOpenAcc aenv (Array (sh, Int) e, Array sh e)
forall aenv e (acc :: * -> * -> *) sh.
Direction
-> Fun aenv (e -> e -> e)
-> Exp aenv e
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array (sh, Int) e, Array sh e)
Scan'    Direction
d Fun aenv (e -> e -> e)
f Exp aenv e
z (Config
-> OpenAcc aenv (Array (sh, Int) e)
-> DelayedOpenAcc aenv (Array (sh, Int) e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array (sh, Int) e)
a)
    Permute Fun aenv (e -> e -> e)
f OpenAcc aenv (Array sh' e)
d Fun aenv (sh -> PrimMaybe sh')
p OpenAcc aenv (Array sh e)
a         -> Fun aenv (e -> e -> e)
-> DelayedOpenAcc aenv (Array sh' e)
-> Fun aenv (sh -> PrimMaybe sh')
-> DelayedOpenAcc aenv (Array sh e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh' e)
forall aenv e (acc :: * -> * -> *) sh' sh.
Fun aenv (e -> e -> e)
-> acc aenv (Array sh' e)
-> Fun aenv (sh -> PrimMaybe sh')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Permute  Fun aenv (e -> e -> e)
f (Config
-> OpenAcc aenv (Array sh' e) -> DelayedOpenAcc aenv (Array sh' e)
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv (Array sh' e)
d) Fun aenv (sh -> PrimMaybe sh')
p (Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array sh e)
a)
    Stencil StencilR sh e stencil
s TypeR e'
t Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
x OpenAcc aenv (Array sh e)
a       -> StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> DelayedOpenAcc aenv (Array sh e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e')
forall sh e stencil e' aenv (acc :: * -> * -> *).
StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Stencil  StencilR sh e stencil
s TypeR e'
t Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
x (Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array sh e)
a)
    Stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a1)
x OpenAcc aenv (Array sh a1)
a Boundary aenv (Array sh b)
y OpenAcc aenv (Array sh b)
b
                            -> StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> DelayedOpenAcc aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> DelayedOpenAcc aenv (Array sh b)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh c)
forall sh a1 stencil1 b stencil2 c aenv (acc :: * -> * -> *).
StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> acc aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> acc aenv (Array sh b)
-> PreOpenAcc acc aenv (Array sh c)
Stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a1)
x (Config
-> OpenAcc aenv (Array sh a1) -> DelayedOpenAcc aenv (Array sh a1)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array sh a1)
a) Boundary aenv (Array sh b)
y (Config
-> OpenAcc aenv (Array sh b) -> DelayedOpenAcc aenv (Array sh b)
forall aenv sh e.
HasCallStack =>
Config
-> OpenAcc aenv (Array sh e) -> DelayedOpenAcc aenv (Array sh e)
delayed Config
config OpenAcc aenv (Array sh b)
b)
    -- Collect s               -> Collect  (cvtS s)

    where
      -- Flatten needless let-binds, which can be introduced by the
      -- conversion to the internal embeddable representation.
      --
      alet :: HasCallStack
           => ALeftHandSide a aenv aenv'
           -> DelayedOpenAcc aenv a
           -> DelayedOpenAcc aenv' b
           -> PreOpenAcc DelayedOpenAcc aenv b
      alet :: forall a aenv aenv' b.
HasCallStack =>
ALeftHandSide a aenv aenv'
-> DelayedOpenAcc aenv a
-> DelayedOpenAcc aenv' b
-> PreOpenAcc DelayedOpenAcc aenv b
alet ALeftHandSide a aenv aenv'
lhs DelayedOpenAcc aenv a
bnd DelayedOpenAcc aenv' b
body
        | Just ArrayVars aenv' b
bodyVars  <- DelayedOpenAcc aenv' b -> Maybe (ArrayVars aenv' b)
forall aenv a. DelayedOpenAcc aenv a -> Maybe (ArrayVars aenv a)
extractDelayedArrayVars DelayedOpenAcc aenv' b
body
        , Just a :~: b
Refl      <- ALeftHandSide a aenv aenv' -> ArrayVars aenv' b -> Maybe (a :~: b)
forall (s :: * -> *) a env1 env2 b.
LeftHandSide s a env1 env2 -> Vars s env2 b -> Maybe (a :~: b)
bindingIsTrivial ALeftHandSide a aenv aenv'
lhs ArrayVars aenv' b
bodyVars
        , Manifest PreOpenAcc DelayedOpenAcc aenv a
x     <- DelayedOpenAcc aenv a
bnd
        = PreOpenAcc DelayedOpenAcc aenv a
PreOpenAcc DelayedOpenAcc aenv b
x
        --
        | Bool
otherwise
        = ALeftHandSide a aenv aenv'
-> DelayedOpenAcc aenv a
-> DelayedOpenAcc aenv' b
-> PreOpenAcc DelayedOpenAcc aenv b
forall bndArrs aenv aenv' (acc :: * -> * -> *) a.
ALeftHandSide bndArrs aenv aenv'
-> acc aenv bndArrs -> acc aenv' a -> PreOpenAcc acc aenv a
Alet ALeftHandSide a aenv aenv'
lhs DelayedOpenAcc aenv a
bnd DelayedOpenAcc aenv' b
body

      -- Eliminate redundant application to an identity function. This
      -- arises in the use of pipe to avoid fusion and force its argument
      -- to be evaluated, i.e.:
      --
      -- > compute :: Acc a -> Acc a
      -- > compute = id >-> id
      --
      apply :: HasCallStack
            => ArraysR b
            -> PreOpenAfun DelayedOpenAcc aenv (a -> b)
            ->             DelayedOpenAcc aenv a
            -> PreOpenAcc  DelayedOpenAcc aenv b
      apply :: forall b aenv a.
HasCallStack =>
ArraysR b
-> PreOpenAfun DelayedOpenAcc aenv (a -> b)
-> DelayedOpenAcc aenv a
-> PreOpenAcc DelayedOpenAcc aenv b
apply ArraysR b
repr PreOpenAfun DelayedOpenAcc aenv (a -> b)
afun DelayedOpenAcc aenv a
x
        | Alam ALeftHandSide a aenv aenv'
lhs (Abody DelayedOpenAcc aenv' t1
body)   <- PreOpenAfun DelayedOpenAcc aenv (a -> b)
afun
        , Just ArrayVars aenv' t1
bodyVars           <- DelayedOpenAcc aenv' t1 -> Maybe (ArrayVars aenv' t1)
forall aenv a. DelayedOpenAcc aenv a -> Maybe (ArrayVars aenv a)
extractDelayedArrayVars DelayedOpenAcc aenv' t1
body
        , Just a :~: t1
Refl               <- ALeftHandSide a aenv aenv'
-> ArrayVars aenv' t1 -> Maybe (a :~: t1)
forall (s :: * -> *) a env1 env2 b.
LeftHandSide s a env1 env2 -> Vars s env2 b -> Maybe (a :~: b)
bindingIsTrivial ALeftHandSide a aenv aenv'
lhs ArrayVars aenv' t1
bodyVars
        , Manifest PreOpenAcc DelayedOpenAcc aenv a
x'             <- DelayedOpenAcc aenv a
x
        = Text
-> PreOpenAcc DelayedOpenAcc aenv b
-> PreOpenAcc DelayedOpenAcc aenv b
forall a. Text -> a -> a
Stats.ruleFired Text
"applyD/identity" PreOpenAcc DelayedOpenAcc aenv b
PreOpenAcc DelayedOpenAcc aenv a
x'
        --
        | Bool
otherwise
        = ArraysR b
-> PreOpenAfun DelayedOpenAcc aenv (a -> b)
-> DelayedOpenAcc aenv a
-> PreOpenAcc DelayedOpenAcc aenv b
forall a (acc :: * -> * -> *) aenv arrs1.
ArraysR a
-> PreOpenAfun acc aenv (arrs1 -> a)
-> acc aenv arrs1
-> PreOpenAcc acc aenv a
Apply ArraysR b
repr PreOpenAfun DelayedOpenAcc aenv (a -> b)
afun DelayedOpenAcc aenv a
x

      cvtAF :: HasCallStack => OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
      cvtAF :: forall aenv f.
HasCallStack =>
OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
cvtAF (Alam ALeftHandSide a aenv aenv'
lhs PreOpenAfun OpenAcc aenv' t1
f) = ALeftHandSide a aenv aenv'
-> PreOpenAfun DelayedOpenAcc aenv' t1
-> PreOpenAfun DelayedOpenAcc aenv (a -> t1)
forall a aenv aenv' (acc :: * -> * -> *) t1.
ALeftHandSide a aenv aenv'
-> PreOpenAfun acc aenv' t1 -> PreOpenAfun acc aenv (a -> t1)
Alam ALeftHandSide a aenv aenv'
lhs (PreOpenAfun OpenAcc aenv' t1 -> PreOpenAfun DelayedOpenAcc aenv' t1
forall aenv f.
HasCallStack =>
OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
cvtAF PreOpenAfun OpenAcc aenv' t1
f)
      cvtAF (Abody OpenAcc aenv f
b)    = DelayedOpenAcc aenv f -> PreOpenAfun DelayedOpenAcc aenv f
forall (acc :: * -> * -> *) aenv t.
acc aenv t -> PreOpenAfun acc aenv t
Abody (Config -> OpenAcc aenv f -> DelayedOpenAcc aenv f
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
manifest Config
config OpenAcc aenv f
b)

      -- cvtS :: PreOpenSeq OpenAcc aenv senv s -> PreOpenSeq DelayedOpenAcc aenv senv s
      -- cvtS = convertOpenSeq config

convertOpenAfun :: HasCallStack => Config -> OpenAfun aenv f -> DelayedOpenAfun aenv f
convertOpenAfun :: forall aenv f.
HasCallStack =>
Config -> OpenAfun aenv f -> DelayedOpenAfun aenv f
convertOpenAfun Config
c (Alam ALeftHandSide a aenv aenv'
lhs PreOpenAfun OpenAcc aenv' t1
f) = ALeftHandSide a aenv aenv'
-> PreOpenAfun DelayedOpenAcc aenv' t1
-> PreOpenAfun DelayedOpenAcc aenv (a -> t1)
forall a aenv aenv' (acc :: * -> * -> *) t1.
ALeftHandSide a aenv aenv'
-> PreOpenAfun acc aenv' t1 -> PreOpenAfun acc aenv (a -> t1)
Alam ALeftHandSide a aenv aenv'
lhs (Config
-> PreOpenAfun OpenAcc aenv' t1
-> PreOpenAfun DelayedOpenAcc aenv' t1
forall aenv f.
HasCallStack =>
Config -> OpenAfun aenv f -> DelayedOpenAfun aenv f
convertOpenAfun Config
c PreOpenAfun OpenAcc aenv' t1
f)
convertOpenAfun Config
c (Abody OpenAcc aenv f
b) = DelayedOpenAcc aenv f -> PreOpenAfun DelayedOpenAcc aenv f
forall (acc :: * -> * -> *) aenv t.
acc aenv t -> PreOpenAfun acc aenv t
Abody (Config -> OpenAcc aenv f -> DelayedOpenAcc aenv f
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> DelayedOpenAcc aenv arrs
convertOpenAcc  Config
c OpenAcc aenv f
b)

{--
convertOpenSeq :: Config -> PreOpenSeq OpenAcc aenv senv a -> PreOpenSeq DelayedOpenAcc aenv senv a
convertOpenSeq config s =
  case s of
    Consumer c          -> Consumer (cvtC c)
    Reify ix            -> Reify ix
    Producer p s'       -> Producer p' (convertOpenSeq config s')
      where
        p' = case p of
               StreamIn arrs     -> StreamIn arrs
               ToSeq slix sh a   -> ToSeq slix sh (delayed config a)
               MapSeq f x        -> MapSeq (cvtAF f) x
               ChunkedMapSeq f x -> ChunkedMapSeq (cvtAF f) x
               ZipWithSeq f x y  -> ZipWithSeq (cvtAF f) x y
               ScanSeq f e x     -> ScanSeq (cvtF f) (cvtE e) x
  where
    cvtC :: Consumer OpenAcc aenv senv a -> Consumer DelayedOpenAcc aenv senv a
    cvtC c =
      case c of
        FoldSeq f e x        -> FoldSeq (cvtF f) (cvtE e) x
        FoldSeqFlatten f a x -> FoldSeqFlatten (cvtAF f) (manifest config a) x
        Stuple t             -> Stuple (cvtCT t)

    cvtCT :: Atuple (Consumer OpenAcc aenv senv) t -> Atuple (Consumer DelayedOpenAcc aenv senv) t
    cvtCT NilAtup        = NilAtup
    cvtCT (SnocAtup t c) = SnocAtup (cvtCT t) (cvtC c)

    cvtAF :: OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f
    cvtAF (Alam f)  = Alam  (cvtAF f)
    cvtAF (Abody b) = Abody (manifest config b)

    cvtE :: OpenExp env aenv t -> DelayedOpenExp env aenv t
    cvtE = convertOpenExp config

    cvtF :: OpenFun env aenv f -> DelayedOpenFun env aenv f
    cvtF (Lam f)  = Lam (cvtF f)
    cvtF (Body b) = Body (cvtE b)
--}


type EmbedAcc acc = forall aenv arrs. acc aenv arrs -> Embed acc aenv arrs
type ElimAcc  acc = forall aenv s t. acc aenv s -> acc (aenv,s) t -> Bool


-- | Apply the fusion transformation to the AST to combine and simplify terms.
-- This converts terms into the internal delayed array representation and merges
-- adjacent producer/producer terms. Using the reduced internal form limits the
-- number of combinations that need to be considered.
--
embedOpenAcc :: HasCallStack => Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
embedOpenAcc :: forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
embedOpenAcc Config
config (OpenAcc PreOpenAcc OpenAcc aenv arrs
pacc) =
  Config
-> MatchAcc OpenAcc
-> EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> PreOpenAcc OpenAcc aenv arrs
-> Embed OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Config
-> MatchAcc OpenAcc
-> EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> PreOpenAcc OpenAcc aenv arrs
-> Embed OpenAcc aenv arrs
embedPreOpenAcc Config
config OpenAcc aenv s -> OpenAcc aenv t -> Maybe (s :~: t)
MatchAcc OpenAcc
matchOpenAcc (Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
embedOpenAcc Config
config) OpenAcc aenv s -> OpenAcc (aenv, s) t -> Bool
ElimAcc OpenAcc
elimOpenAcc PreOpenAcc OpenAcc aenv arrs
pacc
  where
    -- When does the cost of re-computation outweigh that of memory access? For
    -- the moment only do the substitution on a single use of the bound array
    -- into the use site, but it is likely advantageous to be far more
    -- aggressive here.
    --
    -- SEE: [Sharing vs. Fusion]
    --
    elimOpenAcc :: ElimAcc OpenAcc
    elimOpenAcc :: ElimAcc OpenAcc
elimOpenAcc OpenAcc aenv s
_bnd OpenAcc (aenv, s) t
body
      | Bool -> Idx (aenv, s) s -> OpenAcc (aenv, s) t -> Int
UsesOfAcc OpenAcc
count Bool
False Idx (aenv, s) s
forall envt t env. (envt ~ (env, t)) => Idx envt t
ZeroIdx OpenAcc (aenv, s) t
body Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
lIMIT = Bool
True
      | Bool
otherwise                         = Bool
False
      where
        lIMIT :: Int
lIMIT = Int
1

        count :: UsesOfAcc OpenAcc
        count :: UsesOfAcc OpenAcc
count Bool
no Idx aenv s
ix (OpenAcc PreOpenAcc OpenAcc aenv t
pacc) = Bool
-> UsesOfAcc OpenAcc
-> Idx aenv s
-> PreOpenAcc OpenAcc aenv t
-> Int
forall (acc :: * -> * -> *) aenv s t.
Bool -> UsesOfAcc acc -> Idx aenv s -> PreOpenAcc acc aenv t -> Int
usesOfPreAcc Bool
no Bool -> Idx aenv s -> OpenAcc aenv t -> Int
UsesOfAcc OpenAcc
count Idx aenv s
ix PreOpenAcc OpenAcc aenv t
pacc

    matchOpenAcc :: MatchAcc OpenAcc
    matchOpenAcc :: MatchAcc OpenAcc
matchOpenAcc (OpenAcc PreOpenAcc OpenAcc aenv s
pacc1) (OpenAcc PreOpenAcc OpenAcc aenv t
pacc2) =
      MatchAcc OpenAcc
-> PreOpenAcc OpenAcc aenv s
-> PreOpenAcc OpenAcc aenv t
-> Maybe (s :~: t)
forall (acc :: * -> * -> *) aenv s t.
HasArraysR acc =>
MatchAcc acc
-> PreOpenAcc acc aenv s
-> PreOpenAcc acc aenv t
-> Maybe (s :~: t)
matchPreOpenAcc OpenAcc aenv s -> OpenAcc aenv t -> Maybe (s :~: t)
MatchAcc OpenAcc
matchOpenAcc PreOpenAcc OpenAcc aenv s
pacc1 PreOpenAcc OpenAcc aenv t
pacc2


embedPreOpenAcc
    :: HasCallStack
    => Config
    -> MatchAcc   OpenAcc
    -> EmbedAcc   OpenAcc
    -> ElimAcc    OpenAcc
    -> PreOpenAcc OpenAcc aenv arrs
    -> Embed      OpenAcc aenv arrs
embedPreOpenAcc :: forall aenv arrs.
HasCallStack =>
Config
-> MatchAcc OpenAcc
-> EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> PreOpenAcc OpenAcc aenv arrs
-> Embed OpenAcc aenv arrs
embedPreOpenAcc Config
config MatchAcc OpenAcc
matchAcc EmbedAcc OpenAcc
embedAcc ElimAcc OpenAcc
elimAcc PreOpenAcc OpenAcc aenv arrs
pacc
  = Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
unembed
  (Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ case PreOpenAcc OpenAcc aenv arrs
pacc of

    -- Non-fusible terms
    -- -----------------
    --
    -- Solid and semi-solid terms that we generally do not wish to fuse, such
    -- as control flow (|?), array introduction (use, unit), array tupling and
    -- projection, and foreign function operations. Generally we also do not
    -- want to fuse past array let bindings, as this would imply work
    -- duplication. SEE: [Sharing vs. Fusion]
    --
    Alet ALeftHandSide bndArrs aenv aenv'
lhs OpenAcc aenv bndArrs
bnd OpenAcc aenv' arrs
body   -> EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> ALeftHandSide bndArrs aenv aenv'
-> OpenAcc aenv bndArrs
-> OpenAcc aenv' arrs
-> Embed OpenAcc aenv arrs
forall arrs aenv aenv' brrs.
HasCallStack =>
EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> ALeftHandSide arrs aenv aenv'
-> OpenAcc aenv arrs
-> OpenAcc aenv' brrs
-> Embed OpenAcc aenv brrs
aletD OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv s -> OpenAcc (aenv, s) t -> Bool
ElimAcc OpenAcc
elimAcc ALeftHandSide bndArrs aenv aenv'
lhs OpenAcc aenv bndArrs
bnd OpenAcc aenv' arrs
body
    PreOpenAcc OpenAcc aenv arrs
Anil                -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ PreOpenAcc OpenAcc aenv arrs
PreOpenAcc OpenAcc aenv ()
forall (acc :: * -> * -> *) aenv. PreOpenAcc acc aenv ()
Anil
    Acond Exp aenv PrimBool
p OpenAcc aenv arrs
at OpenAcc aenv arrs
ae       -> MatchAcc OpenAcc
-> EmbedAcc OpenAcc
-> Exp aenv PrimBool
-> OpenAcc aenv arrs
-> OpenAcc aenv arrs
-> Embed OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
MatchAcc OpenAcc
-> EmbedAcc OpenAcc
-> Exp aenv PrimBool
-> OpenAcc aenv arrs
-> OpenAcc aenv arrs
-> Embed OpenAcc aenv arrs
acondD OpenAcc aenv s -> OpenAcc aenv t -> Maybe (s :~: t)
MatchAcc OpenAcc
matchAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc (Exp aenv PrimBool -> Exp aenv PrimBool
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv PrimBool
p) OpenAcc aenv arrs
at OpenAcc aenv arrs
ae
    Apply ArraysR arrs
aR PreOpenAfun OpenAcc aenv (arrs1 -> arrs)
f OpenAcc aenv arrs1
a        -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ ArraysR arrs
-> PreOpenAfun OpenAcc aenv (arrs1 -> arrs)
-> OpenAcc aenv arrs1
-> PreOpenAcc OpenAcc aenv arrs
forall a (acc :: * -> * -> *) aenv arrs1.
ArraysR a
-> PreOpenAfun acc aenv (arrs1 -> a)
-> acc aenv arrs1
-> PreOpenAcc acc aenv a
Apply ArraysR arrs
aR (PreOpenAfun OpenAcc aenv (arrs1 -> arrs)
-> PreOpenAfun OpenAcc aenv (arrs1 -> arrs)
forall aenv' f.
HasCallStack =>
PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
cvtAF PreOpenAfun OpenAcc aenv (arrs1 -> arrs)
f) (OpenAcc aenv arrs1 -> OpenAcc aenv arrs1
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv arrs1
a)
    Awhile PreOpenAfun OpenAcc aenv (arrs -> Scalar PrimBool)
p PreOpenAfun OpenAcc aenv (arrs -> arrs)
f OpenAcc aenv arrs
a        -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ PreOpenAfun OpenAcc aenv (arrs -> Scalar PrimBool)
-> PreOpenAfun OpenAcc aenv (arrs -> arrs)
-> OpenAcc aenv arrs
-> PreOpenAcc OpenAcc aenv arrs
forall (acc :: * -> * -> *) aenv a.
PreOpenAfun acc aenv (a -> Scalar PrimBool)
-> PreOpenAfun acc aenv (a -> a)
-> acc aenv a
-> PreOpenAcc acc aenv a
Awhile (PreOpenAfun OpenAcc aenv (arrs -> Scalar PrimBool)
-> PreOpenAfun OpenAcc aenv (arrs -> Scalar PrimBool)
forall aenv' f.
HasCallStack =>
PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
cvtAF PreOpenAfun OpenAcc aenv (arrs -> Scalar PrimBool)
p) (PreOpenAfun OpenAcc aenv (arrs -> arrs)
-> PreOpenAfun OpenAcc aenv (arrs -> arrs)
forall aenv' f.
HasCallStack =>
PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
cvtAF PreOpenAfun OpenAcc aenv (arrs -> arrs)
f) (OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv arrs
a)
    Apair OpenAcc aenv as
a1 OpenAcc aenv bs
a2         -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ OpenAcc aenv as
-> OpenAcc aenv bs -> PreOpenAcc OpenAcc aenv (as, bs)
forall (acc :: * -> * -> *) aenv as bs.
acc aenv as -> acc aenv bs -> PreOpenAcc acc aenv (as, bs)
Apair (OpenAcc aenv as -> OpenAcc aenv as
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv as
a1) (OpenAcc aenv bs -> OpenAcc aenv bs
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv bs
a2)
    Atrace Message arrs1
msg OpenAcc aenv arrs1
a1 OpenAcc aenv arrs
a2    -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ Message arrs1
-> OpenAcc aenv arrs1
-> OpenAcc aenv arrs
-> PreOpenAcc OpenAcc aenv arrs
forall arrs1 (acc :: * -> * -> *) aenv a.
Message arrs1
-> acc aenv arrs1 -> acc aenv a -> PreOpenAcc acc aenv a
Atrace Message arrs1
msg (OpenAcc aenv arrs1 -> OpenAcc aenv arrs1
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv arrs1
a1) (OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv arrs
a2)
    Aforeign ArraysR arrs
aR asm (as -> arrs)
ff PreAfun OpenAcc (as -> arrs)
f OpenAcc aenv as
a  -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ ArraysR arrs
-> asm (as -> arrs)
-> PreAfun OpenAcc (as -> arrs)
-> OpenAcc aenv as
-> PreOpenAcc OpenAcc aenv arrs
forall (asm :: * -> *) a as (acc :: * -> * -> *) aenv.
Foreign asm =>
ArraysR a
-> asm (as -> a)
-> PreAfun acc (as -> a)
-> acc aenv as
-> PreOpenAcc acc aenv a
Aforeign ArraysR arrs
aR asm (as -> arrs)
ff (PreAfun OpenAcc (as -> arrs) -> PreAfun OpenAcc (as -> arrs)
forall aenv' f.
HasCallStack =>
PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
cvtAF PreAfun OpenAcc (as -> arrs)
f) (OpenAcc aenv as -> OpenAcc aenv as
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv as
a)
    -- Collect s           -> collectD s

    -- Array injection
    Avar ArrayVar aenv (Array sh e)
v              -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ ArrayVar aenv (Array sh e) -> PreOpenAcc OpenAcc aenv (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar ArrayVar aenv (Array sh e)
v
    Use ArrayR (Array sh e)
aR Array sh e
a            -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e)
-> Array sh e -> PreOpenAcc OpenAcc aenv (Array sh e)
forall sh e (acc :: * -> * -> *) aenv.
ArrayR (Array sh e)
-> Array sh e -> PreOpenAcc acc aenv (Array sh e)
Use ArrayR (Array sh e)
aR Array sh e
a
    Unit TypeR e
t Exp aenv e
e            -> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ TypeR e -> Exp aenv e -> PreOpenAcc OpenAcc aenv (Array () e)
forall e aenv (acc :: * -> * -> *).
TypeR e -> Exp aenv e -> PreOpenAcc acc aenv (Array () e)
Unit TypeR e
t (Exp aenv e -> Exp aenv e
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv e
e)

    -- Producers
    -- ---------
    --
    -- The class of operations that given a set of zero or more input arrays,
    -- produce a _single_ element for the output array by manipulating a
    -- _single_ element from each input array. These can be further classified
    -- as value (map, zipWith) or index space (backpermute, slice, replicate)
    -- transformations.
    --
    -- The critical feature is that each element of the output is produced
    -- independently of all others, and so we can aggressively fuse arbitrary
    -- sequences of these operations.
    --
    Generate ArrayR (Array sh e)
aR Exp aenv sh
sh Fun aenv (sh -> e)
f        -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Embed OpenAcc aenv (Array sh e)
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Embed OpenAcc aenv (Array sh e)
generateD ArrayR (Array sh e)
aR (Exp aenv sh -> Exp aenv sh
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv sh
sh) (Fun aenv (sh -> e) -> Fun aenv (sh -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (sh -> e)
f)
    Map TypeR e'
t Fun aenv (e -> e')
f OpenAcc aenv (Array sh e)
a               -> TypeR e'
-> Fun aenv (e -> e')
-> Embed OpenAcc aenv (Array sh e)
-> Embed OpenAcc aenv (Array sh e')
forall b aenv a sh.
HasCallStack =>
TypeR b
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh b)
mapD TypeR e'
t (Fun aenv (e -> e') -> Fun aenv (e -> e')
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (e -> e')
f) (OpenAcc aenv (Array sh e) -> Embed OpenAcc aenv (Array sh e)
EmbedAcc OpenAcc
embedAcc OpenAcc aenv (Array sh e)
a)
    ZipWith TypeR e3
t Fun aenv (e1 -> e2 -> e3)
f OpenAcc aenv (Array sh e1)
a OpenAcc aenv (Array sh e2)
b         -> (forall r s aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' (Array sh e1)
 -> Cunctation s aenv' (Array sh e2)
 -> Cunctation D aenv' arrs)
-> OpenAcc aenv (Array sh e1)
-> OpenAcc aenv (Array sh e2)
-> Embed OpenAcc aenv arrs
forall aenv as bs cs.
HasCallStack =>
(forall r s aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as
 -> Cunctation s aenv' bs
 -> Cunctation D aenv' cs)
-> OpenAcc aenv as -> OpenAcc aenv bs -> Embed OpenAcc aenv cs
fuse2 ((OpenFun () aenv' (e1 -> e2 -> e3)
 -> Cunctation r aenv' (Array sh e1)
 -> Cunctation s aenv' (Array sh e2)
 -> Cunctation D aenv' arrs)
-> Fun aenv (e1 -> e2 -> e3)
-> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e1)
-> Cunctation s aenv' (Array sh e2)
-> Cunctation D aenv' arrs
forall (f :: * -> * -> *) env' a b env.
(HasCallStack, Sink f) =>
(f env' a -> b) -> f env a -> Extend ArrayR OpenAcc env env' -> b
into (TypeR e3
-> OpenFun () aenv' (e1 -> e2 -> e3)
-> Cunctation r aenv' (Array sh e1)
-> Cunctation s aenv' (Array sh e2)
-> Cunctation D aenv' (Array sh e3)
forall c aenv a b r sh s.
HasCallStack =>
TypeR c
-> Fun aenv (a -> b -> c)
-> Cunctation r aenv (Array sh a)
-> Cunctation s aenv (Array sh b)
-> Cunctation D aenv (Array sh c)
zipWithD TypeR e3
t) (Fun aenv (e1 -> e2 -> e3) -> Fun aenv (e1 -> e2 -> e3)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (e1 -> e2 -> e3)
f)) OpenAcc aenv (Array sh e1)
a OpenAcc aenv (Array sh e2)
b
    Transform ArrayR (Array sh' b)
aR Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a1 -> b)
f OpenAcc aenv (Array sh a1)
a   -> ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> Embed OpenAcc aenv (Array sh a1)
-> Embed OpenAcc aenv (Array sh' b)
forall sh' b aenv sh a.
HasCallStack =>
ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh' b)
transformD ArrayR (Array sh' b)
aR (Exp aenv sh' -> Exp aenv sh'
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv sh'
sh) (Fun aenv (sh' -> sh) -> Fun aenv (sh' -> sh)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (sh' -> sh)
p) (Fun aenv (a1 -> b) -> Fun aenv (a1 -> b)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (a1 -> b)
f) (OpenAcc aenv (Array sh a1) -> Embed OpenAcc aenv (Array sh a1)
EmbedAcc OpenAcc
embedAcc OpenAcc aenv (Array sh a1)
a)
    Backpermute ShapeR sh'
slr Exp aenv sh'
sl Fun aenv (sh' -> sh)
p OpenAcc aenv (Array sh e)
a  -> (forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' (Array sh e) -> Cunctation D aenv' arrs)
-> OpenAcc aenv (Array sh e) -> Embed OpenAcc aenv arrs
forall aenv as bs.
HasCallStack =>
(forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as -> Cunctation D aenv' bs)
-> OpenAcc aenv as -> Embed OpenAcc aenv bs
fuse ((OpenExp () aenv' sh'
 -> OpenFun () aenv' (sh' -> sh)
 -> Cunctation r aenv' (Array sh e)
 -> Cunctation D aenv' arrs)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> Cunctation D aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env.
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> f2 env' b -> c)
-> f1 env a -> f2 env b -> Extend ArrayR OpenAcc env env' -> c
into2 (ShapeR sh'
-> OpenExp () aenv' sh'
-> OpenFun () aenv' (sh' -> sh)
-> Cunctation r aenv' (Array sh e)
-> Cunctation D aenv' (Array sh' e)
forall sh' aenv sh r e.
HasCallStack =>
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
backpermuteD ShapeR sh'
slr) (Exp aenv sh' -> Exp aenv sh'
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv sh'
sl) (Fun aenv (sh' -> sh) -> Fun aenv (sh' -> sh)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (sh' -> sh)
p)) OpenAcc aenv (Array sh e)
a
    Slice SliceIndex slix sl co sh
slix OpenAcc aenv (Array sh e)
a Exp aenv slix
sl         -> (forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' (Array sh e) -> Cunctation D aenv' arrs)
-> OpenAcc aenv (Array sh e) -> Embed OpenAcc aenv arrs
forall aenv as bs.
HasCallStack =>
(forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as -> Cunctation D aenv' bs)
-> OpenAcc aenv as -> Embed OpenAcc aenv bs
fuse ((OpenExp () aenv' slix
 -> Cunctation r aenv' (Array sh e) -> Cunctation D aenv' arrs)
-> Exp aenv slix
-> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> Cunctation D aenv' arrs
forall (f :: * -> * -> *) env' a b env.
(HasCallStack, Sink f) =>
(f env' a -> b) -> f env a -> Extend ArrayR OpenAcc env env' -> b
into  (SliceIndex slix sl co sh
-> OpenExp () aenv' slix
-> Cunctation r aenv' (Array sh e)
-> Cunctation D aenv' (Array sl e)
forall slix sl co sh aenv r e.
HasCallStack =>
SliceIndex slix sl co sh
-> Exp aenv slix
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sl e)
sliceD SliceIndex slix sl co sh
slix)      (Exp aenv slix -> Exp aenv slix
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv slix
sl)) OpenAcc aenv (Array sh e)
a
    Replicate SliceIndex slix sl co sh
slix Exp aenv slix
sh OpenAcc aenv (Array sl e)
a     -> (forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' (Array sl e) -> Cunctation D aenv' arrs)
-> OpenAcc aenv (Array sl e) -> Embed OpenAcc aenv arrs
forall aenv as bs.
HasCallStack =>
(forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as -> Cunctation D aenv' bs)
-> OpenAcc aenv as -> Embed OpenAcc aenv bs
fuse ((OpenExp () aenv' slix
 -> Cunctation r aenv' (Array sl e) -> Cunctation D aenv' arrs)
-> Exp aenv slix
-> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sl e)
-> Cunctation D aenv' arrs
forall (f :: * -> * -> *) env' a b env.
(HasCallStack, Sink f) =>
(f env' a -> b) -> f env a -> Extend ArrayR OpenAcc env env' -> b
into  (SliceIndex slix sl co sh
-> OpenExp () aenv' slix
-> Cunctation r aenv' (Array sl e)
-> Cunctation D aenv' (Array sh e)
forall slix sl co sh aenv r e.
HasCallStack =>
SliceIndex slix sl co sh
-> Exp aenv slix
-> Cunctation r aenv (Array sl e)
-> Cunctation D aenv (Array sh e)
replicateD SliceIndex slix sl co sh
slix)  (Exp aenv slix -> Exp aenv slix
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv slix
sh)) OpenAcc aenv (Array sl e)
a
    Reshape ShapeR sh
slr Exp aenv sh
sl OpenAcc aenv (Array sh' e)
a        -> ShapeR sh
-> Embed OpenAcc aenv (Array sh' e)
-> Exp aenv sh
-> Embed OpenAcc aenv (Array sh e)
forall sl aenv sh e.
HasCallStack =>
ShapeR sl
-> Embed OpenAcc aenv (Array sh e)
-> Exp aenv sl
-> Embed OpenAcc aenv (Array sl e)
reshapeD ShapeR sh
slr (OpenAcc aenv (Array sh' e) -> Embed OpenAcc aenv (Array sh' e)
EmbedAcc OpenAcc
embedAcc OpenAcc aenv (Array sh' e)
a) (Exp aenv sh -> Exp aenv sh
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv sh
sl)

    -- Consumers
    -- ---------
    --
    -- Operations where each element of the output array depends on multiple
    -- elements of the input array. To implement these operations efficiently in
    -- parallel, we need to know how elements of the array depend on each other:
    -- a parallel scan is implemented very differently from a parallel fold, for
    -- example.
    --
    -- In order to avoid obfuscating this crucial information required for
    -- parallel implementation, fusion is separated into to phases:
    -- producer/producer, implemented above, and consumer/producer, which is
    -- implemented below. This will place producers adjacent to the consumer
    -- node, so that the producer can be directly embedded into the consumer
    -- during the code generation phase.
    --
    Fold Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a                  -> ArraysR arrs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' (Array (sh, Int) e)
    -> PreOpenAcc OpenAcc aenv' arrs)
-> OpenAcc aenv (Array (sh, Int) e)
-> Embed OpenAcc aenv arrs
forall bs aenv as.
HasCallStack =>
ArraysR bs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs)
-> OpenAcc aenv as
-> Embed OpenAcc aenv bs
embed  ArraysR arrs
aR ((OpenFun () aenv' (e -> e -> e)
 -> Maybe (OpenExp () aenv' e)
 -> OpenAcc aenv' (Array (sh, Int) e)
 -> PreOpenAcc OpenAcc aenv' arrs)
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env
       (acc :: * -> * -> *).
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> Maybe (f2 env' b) -> c)
-> f1 env a -> Maybe (f2 env b) -> Extend ArrayR acc env env' -> c
into2M OpenFun () aenv' (e -> e -> e)
-> Maybe (OpenExp () aenv' e)
-> OpenAcc aenv' (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv' arrs
OpenFun () aenv' (e -> e -> e)
-> Maybe (OpenExp () aenv' e)
-> OpenAcc aenv' (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv' (Array sh e)
forall aenv e (acc :: * -> * -> *) sh.
Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array sh e)
Fold               (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE (Exp aenv e -> Exp aenv e)
-> Maybe (Exp aenv e) -> Maybe (Exp aenv e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Exp aenv e)
z)) OpenAcc aenv (Array (sh, Int) e)
a
    FoldSeg IntegralType i
i Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a OpenAcc aenv (Segments i)
s           -> ArraysR arrs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' (Array (sh, Int) e)
    -> OpenAcc aenv' (Segments i)
    -> PreOpenAcc OpenAcc aenv' arrs)
-> OpenAcc aenv (Array (sh, Int) e)
-> OpenAcc aenv (Segments i)
-> Embed OpenAcc aenv arrs
forall cs aenv as bs.
HasCallStack =>
ArraysR cs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as
    -> OpenAcc aenv' bs
    -> PreOpenAcc OpenAcc aenv' cs)
-> OpenAcc aenv as
-> OpenAcc aenv bs
-> Embed OpenAcc aenv cs
embed2 ArraysR arrs
aR ((OpenFun () aenv' (e -> e -> e)
 -> Maybe (OpenExp () aenv' e)
 -> OpenAcc aenv' (Array (sh, Int) e)
 -> OpenAcc aenv' (Segments i)
 -> PreOpenAcc OpenAcc aenv' arrs)
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array (sh, Int) e)
-> OpenAcc aenv' (Segments i)
-> PreOpenAcc OpenAcc aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env
       (acc :: * -> * -> *).
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> Maybe (f2 env' b) -> c)
-> f1 env a -> Maybe (f2 env b) -> Extend ArrayR acc env env' -> c
into2M (IntegralType i
-> OpenFun () aenv' (e -> e -> e)
-> Maybe (OpenExp () aenv' e)
-> OpenAcc aenv' (Array (sh, Int) e)
-> OpenAcc aenv' (Segments i)
-> PreOpenAcc OpenAcc aenv' (Array (sh, Int) e)
forall i aenv e (acc :: * -> * -> *) sh.
IntegralType i
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> acc aenv (Segments i)
-> PreOpenAcc acc aenv (Array (sh, Int) e)
FoldSeg IntegralType i
i)        (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE (Exp aenv e -> Exp aenv e)
-> Maybe (Exp aenv e) -> Maybe (Exp aenv e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Exp aenv e)
z)) OpenAcc aenv (Array (sh, Int) e)
a OpenAcc aenv (Segments i)
s
    Scan  Direction
d Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a               -> ArraysR arrs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' (Array (sh, Int) e)
    -> PreOpenAcc OpenAcc aenv' arrs)
-> OpenAcc aenv (Array (sh, Int) e)
-> Embed OpenAcc aenv arrs
forall bs aenv as.
HasCallStack =>
ArraysR bs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs)
-> OpenAcc aenv as
-> Embed OpenAcc aenv bs
embed  ArraysR arrs
aR ((OpenFun () aenv' (e -> e -> e)
 -> Maybe (OpenExp () aenv' e)
 -> OpenAcc aenv' (Array (sh, Int) e)
 -> PreOpenAcc OpenAcc aenv' arrs)
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env
       (acc :: * -> * -> *).
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> Maybe (f2 env' b) -> c)
-> f1 env a -> Maybe (f2 env b) -> Extend ArrayR acc env env' -> c
into2M (Direction
-> OpenFun () aenv' (e -> e -> e)
-> Maybe (OpenExp () aenv' e)
-> OpenAcc aenv' (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv' (Array (sh, Int) e)
forall aenv e (acc :: * -> * -> *) sh.
Direction
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array (sh, Int) e)
Scan  Direction
d)          (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE (Exp aenv e -> Exp aenv e)
-> Maybe (Exp aenv e) -> Maybe (Exp aenv e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Exp aenv e)
z)) OpenAcc aenv (Array (sh, Int) e)
a
    Scan' Direction
d Fun aenv (e -> e -> e)
f Exp aenv e
z OpenAcc aenv (Array (sh, Int) e)
a               -> ArraysR arrs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' (Array (sh, Int) e)
    -> PreOpenAcc OpenAcc aenv' arrs)
-> OpenAcc aenv (Array (sh, Int) e)
-> Embed OpenAcc aenv arrs
forall bs aenv as.
HasCallStack =>
ArraysR bs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs)
-> OpenAcc aenv as
-> Embed OpenAcc aenv bs
embed  ArraysR arrs
aR ((OpenFun () aenv' (e -> e -> e)
 -> OpenExp () aenv' e
 -> OpenAcc aenv' (Array (sh, Int) e)
 -> PreOpenAcc OpenAcc aenv' arrs)
-> Fun aenv (e -> e -> e)
-> Exp aenv e
-> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env.
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> f2 env' b -> c)
-> f1 env a -> f2 env b -> Extend ArrayR OpenAcc env env' -> c
into2  (Direction
-> OpenFun () aenv' (e -> e -> e)
-> OpenExp () aenv' e
-> OpenAcc aenv' (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv' (Array (sh, Int) e, Array sh e)
forall aenv e (acc :: * -> * -> *) sh.
Direction
-> Fun aenv (e -> e -> e)
-> Exp aenv e
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array (sh, Int) e, Array sh e)
Scan' Direction
d)          (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE Exp aenv e
z)) OpenAcc aenv (Array (sh, Int) e)
a
    Permute Fun aenv (e -> e -> e)
f OpenAcc aenv (Array sh' e)
d Fun aenv (sh -> PrimMaybe sh')
p OpenAcc aenv (Array sh e)
a             -> ArraysR arrs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' (Array sh' e)
    -> OpenAcc aenv' (Array sh e)
    -> PreOpenAcc OpenAcc aenv' arrs)
-> OpenAcc aenv (Array sh' e)
-> OpenAcc aenv (Array sh e)
-> Embed OpenAcc aenv arrs
forall cs aenv as bs.
HasCallStack =>
ArraysR cs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as
    -> OpenAcc aenv' bs
    -> PreOpenAcc OpenAcc aenv' cs)
-> OpenAcc aenv as
-> OpenAcc aenv bs
-> Embed OpenAcc aenv cs
embed2 ArraysR arrs
aR ((OpenFun () aenv' (e -> e -> e)
 -> OpenFun () aenv' (sh -> PrimMaybe sh')
 -> OpenAcc aenv' (Array sh' e)
 -> OpenAcc aenv' (Array sh e)
 -> PreOpenAcc OpenAcc aenv' arrs)
-> Fun aenv (e -> e -> e)
-> Fun aenv (sh -> PrimMaybe sh')
-> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array sh' e)
-> OpenAcc aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env.
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> f2 env' b -> c)
-> f1 env a -> f2 env b -> Extend ArrayR OpenAcc env env' -> c
into2  OpenFun () aenv' (e -> e -> e)
-> OpenFun () aenv' (sh -> PrimMaybe sh')
-> OpenAcc aenv' (Array sh' e)
-> OpenAcc aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' arrs
OpenFun () aenv' (e -> e -> e)
-> OpenFun () aenv' (sh -> PrimMaybe sh')
-> OpenAcc aenv' (Array sh' e)
-> OpenAcc aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sh' e)
forall {aenv} {e} {sh} {sh'} {acc :: * -> * -> *}.
Fun aenv (e -> e -> e)
-> Fun aenv (sh -> PrimMaybe sh')
-> acc aenv (Array sh' e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
permute            (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (e -> e -> e)
f) (Fun aenv (sh -> PrimMaybe sh') -> Fun aenv (sh -> PrimMaybe sh')
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (sh -> PrimMaybe sh')
p)) OpenAcc aenv (Array sh' e)
d OpenAcc aenv (Array sh e)
a
    Stencil StencilR sh e stencil
s TypeR e'
t Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
x OpenAcc aenv (Array sh e)
a           -> ArraysR arrs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' (Array sh e) -> PreOpenAcc OpenAcc aenv' arrs)
-> OpenAcc aenv (Array sh e)
-> Embed OpenAcc aenv arrs
forall bs aenv as.
HasCallStack =>
ArraysR bs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs)
-> OpenAcc aenv as
-> Embed OpenAcc aenv bs
embed  ArraysR arrs
aR ((OpenFun () aenv' (stencil -> e')
 -> Boundary aenv' (Array sh e)
 -> OpenAcc aenv' (Array sh e)
 -> PreOpenAcc OpenAcc aenv' arrs)
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env.
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> f2 env' b -> c)
-> f1 env a -> f2 env b -> Extend ArrayR OpenAcc env env' -> c
into2  (StencilR sh e stencil
-> TypeR e'
-> OpenFun () aenv' (stencil -> e')
-> Boundary aenv' (Array sh e)
-> OpenAcc aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sh e')
forall sh e stencil e' aenv (acc :: * -> * -> *).
StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
stencil1 StencilR sh e stencil
s TypeR e'
t)     (Fun aenv (stencil -> e') -> Fun aenv (stencil -> e')
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (stencil -> e')
f) (Boundary aenv (Array sh e) -> Boundary aenv (Array sh e)
forall aenv' t.
HasCallStack =>
Boundary aenv' t -> Boundary aenv' t
cvtB Boundary aenv (Array sh e)
x)) OpenAcc aenv (Array sh e)
a
    Stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a1)
x OpenAcc aenv (Array sh a1)
a Boundary aenv (Array sh b)
y OpenAcc aenv (Array sh b)
b  -> ArraysR arrs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' (Array sh a1)
    -> OpenAcc aenv' (Array sh b)
    -> PreOpenAcc OpenAcc aenv' arrs)
-> OpenAcc aenv (Array sh a1)
-> OpenAcc aenv (Array sh b)
-> Embed OpenAcc aenv arrs
forall cs aenv as bs.
HasCallStack =>
ArraysR cs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as
    -> OpenAcc aenv' bs
    -> PreOpenAcc OpenAcc aenv' cs)
-> OpenAcc aenv as
-> OpenAcc aenv bs
-> Embed OpenAcc aenv cs
embed2 ArraysR arrs
aR ((OpenFun () aenv' (stencil1 -> stencil2 -> c)
 -> Boundary aenv' (Array sh a1)
 -> Boundary aenv' (Array sh b)
 -> OpenAcc aenv' (Array sh a1)
 -> OpenAcc aenv' (Array sh b)
 -> PreOpenAcc OpenAcc aenv' arrs)
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array sh a1)
-> OpenAcc aenv' (Array sh b)
-> PreOpenAcc OpenAcc aenv' arrs
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) (f3 :: * -> * -> *)
       env' a b c d env.
(HasCallStack, Sink f1, Sink f2, Sink f3) =>
(f1 env' a -> f2 env' b -> f3 env' c -> d)
-> f1 env a
-> f2 env b
-> f3 env c
-> Extend ArrayR OpenAcc env env'
-> d
into3  (StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> OpenFun () aenv' (stencil1 -> stencil2 -> c)
-> Boundary aenv' (Array sh a1)
-> Boundary aenv' (Array sh b)
-> OpenAcc aenv' (Array sh a1)
-> OpenAcc aenv' (Array sh b)
-> PreOpenAcc OpenAcc aenv' (Array sh c)
forall {sh} {a1} {stencil1} {b} {stencil2} {c} {aenv}
       {acc :: * -> * -> *}.
StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> acc aenv (Array sh a1)
-> acc aenv (Array sh b)
-> PreOpenAcc acc aenv (Array sh c)
stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t) (Fun aenv (stencil1 -> stencil2 -> c)
-> Fun aenv (stencil1 -> stencil2 -> c)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv (stencil1 -> stencil2 -> c)
f) (Boundary aenv (Array sh a1) -> Boundary aenv (Array sh a1)
forall aenv' t.
HasCallStack =>
Boundary aenv' t -> Boundary aenv' t
cvtB Boundary aenv (Array sh a1)
x) (Boundary aenv (Array sh b) -> Boundary aenv (Array sh b)
forall aenv' t.
HasCallStack =>
Boundary aenv' t -> Boundary aenv' t
cvtB Boundary aenv (Array sh b)
y)) OpenAcc aenv (Array sh a1)
a OpenAcc aenv (Array sh b)
b

  where
    aR :: ArraysR arrs
aR = PreOpenAcc OpenAcc aenv arrs -> ArraysR arrs
forall aenv a. PreOpenAcc OpenAcc aenv a -> ArraysR a
forall (f :: * -> * -> *) aenv a.
HasArraysR f =>
f aenv a -> ArraysR a
arraysR PreOpenAcc OpenAcc aenv arrs
pacc

    -- If fusion is not enabled, force terms to the manifest representation
    --
    unembed :: HasCallStack => Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
    unembed :: forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
unembed Embed OpenAcc aenv arrs
x
      | Flag
array_fusion Flag -> BitSet Word32 Flag -> Bool
forall a c. (Enum a, Bits c) => a -> BitSet c a -> Bool
`member` Config -> BitSet Word32 Flag
options Config
config = Embed OpenAcc aenv arrs
x
      | Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' arrs
cc <- Embed OpenAcc aenv arrs
x
      , PreOpenAcc OpenAcc aenv' arrs
pacc         <- Cunctation r aenv' arrs -> PreOpenAcc OpenAcc aenv' arrs
forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute Cunctation r aenv' arrs
cc
      = case ExtractAcc OpenAcc
-> PreOpenAcc OpenAcc aenv' arrs -> Maybe (ArrayVars aenv' arrs)
forall (acc :: * -> * -> *) aenv a.
ExtractAcc acc -> PreOpenAcc acc aenv a -> Maybe (ArrayVars aenv a)
avarsOut OpenAcc env t -> Maybe (PreOpenAcc OpenAcc env t)
ExtractAcc OpenAcc
extractOpenAcc PreOpenAcc OpenAcc aenv' arrs
pacc of
          Just ArrayVars aenv' arrs
vars -> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation M aenv' arrs -> Embed OpenAcc aenv arrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env (Cunctation M aenv' arrs -> Embed OpenAcc aenv arrs)
-> Cunctation M aenv' arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ ArrayVars aenv' arrs -> Cunctation M aenv' arrs
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done ArrayVars aenv' arrs
vars
          Maybe (ArrayVars aenv' arrs)
_
            | DeclareVars LeftHandSide ArrayR arrs aenv' env'
lhs aenv' :> env'
_ forall env''. (env' :> env'') -> Vars ArrayR env'' arrs
value <- TupR ArrayR arrs -> DeclareVars ArrayR arrs aenv'
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (PreOpenAcc OpenAcc aenv' arrs -> TupR ArrayR arrs
forall aenv a. PreOpenAcc OpenAcc aenv a -> ArraysR a
forall (f :: * -> * -> *) aenv a.
HasArraysR f =>
f aenv a -> ArraysR a
arraysR PreOpenAcc OpenAcc aenv' arrs
pacc)
              -> Extend ArrayR OpenAcc aenv env'
-> Cunctation M env' arrs -> Embed OpenAcc aenv arrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
-> LeftHandSide ArrayR arrs aenv' env'
-> OpenAcc aenv' arrs
-> Extend ArrayR OpenAcc aenv env'
forall (s :: * -> *) (f :: * -> * -> *) env env'1 t env'.
Extend s f env env'1
-> LeftHandSide s t env'1 env' -> f env'1 t -> Extend s f env env'
PushEnv Extend ArrayR OpenAcc aenv aenv'
env LeftHandSide ArrayR arrs aenv' env'
lhs (OpenAcc aenv' arrs -> Extend ArrayR OpenAcc aenv env')
-> OpenAcc aenv' arrs -> Extend ArrayR OpenAcc aenv env'
forall a b. (a -> b) -> a -> b
$ PreOpenAcc OpenAcc aenv' arrs -> OpenAcc aenv' arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc PreOpenAcc OpenAcc aenv' arrs
pacc) (Cunctation M env' arrs -> Embed OpenAcc aenv arrs)
-> Cunctation M env' arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ ArrayVars env' arrs -> Cunctation M env' arrs
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done (ArrayVars env' arrs -> Cunctation M env' arrs)
-> ArrayVars env' arrs -> Cunctation M env' arrs
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ArrayVars env' arrs
forall env''. (env' :> env'') -> Vars ArrayR env'' arrs
value env' :> env'
forall env. env :> env
weakenId

    cvtA :: HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
    cvtA :: forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA = Embed OpenAcc aenv' a -> OpenAcc aenv' a
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Embed OpenAcc aenv' a -> OpenAcc aenv' a)
-> (OpenAcc aenv' a -> Embed OpenAcc aenv' a)
-> OpenAcc aenv' a
-> OpenAcc aenv' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenAcc aenv' a -> Embed OpenAcc aenv' a
EmbedAcc OpenAcc
embedAcc

    cvtAF :: HasCallStack => PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
    cvtAF :: forall aenv' f.
HasCallStack =>
PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
cvtAF (Alam ALeftHandSide a aenv' aenv'
lhs PreOpenAfun OpenAcc aenv' t1
f) = ALeftHandSide a aenv' aenv'
-> PreOpenAfun OpenAcc aenv' t1
-> PreOpenAfun OpenAcc aenv' (a -> t1)
forall a aenv aenv' (acc :: * -> * -> *) t1.
ALeftHandSide a aenv aenv'
-> PreOpenAfun acc aenv' t1 -> PreOpenAfun acc aenv (a -> t1)
Alam ALeftHandSide a aenv' aenv'
lhs (PreOpenAfun OpenAcc aenv' t1 -> PreOpenAfun OpenAcc aenv' t1
forall aenv' f.
HasCallStack =>
PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
cvtAF PreOpenAfun OpenAcc aenv' t1
f)
    cvtAF (Abody OpenAcc aenv' f
a)    = OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f
forall (acc :: * -> * -> *) aenv t.
acc aenv t -> PreOpenAfun acc aenv t
Abody (OpenAcc aenv' f -> OpenAcc aenv' f
forall aenv' a. HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a
cvtA OpenAcc aenv' f
a)

    -- Helpers to shuffle the order of arguments to a constructor
    --
    permute :: Fun aenv (e -> e -> e)
-> Fun aenv (sh -> PrimMaybe sh')
-> acc aenv (Array sh' e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
permute Fun aenv (e -> e -> e)
f Fun aenv (sh -> PrimMaybe sh')
p acc aenv (Array sh' e)
d acc aenv (Array sh e)
a     = Fun aenv (e -> e -> e)
-> acc aenv (Array sh' e)
-> Fun aenv (sh -> PrimMaybe sh')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
forall aenv e (acc :: * -> * -> *) sh' sh.
Fun aenv (e -> e -> e)
-> acc aenv (Array sh' e)
-> Fun aenv (sh -> PrimMaybe sh')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Permute Fun aenv (e -> e -> e)
f acc aenv (Array sh' e)
d Fun aenv (sh -> PrimMaybe sh')
p acc aenv (Array sh e)
a

    -- NOTE: [Stencil fusion]
    --
    -- We allow stencils to delay their argument arrays with no special
    -- considerations. This means that the delayed function will be evaluated
    -- _at every element_ of the stencil pattern. We should do some analysis of
    -- when this duplication is beneficial (keeping in mind that the stencil
    -- implementations themselves may share neighbouring elements).
    --
    stencil1 :: StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
stencil1 StencilR sh e stencil
s TypeR e'
t Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
x acc aenv (Array sh e)
a          = StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
forall sh e stencil e' aenv (acc :: * -> * -> *).
StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Stencil  StencilR sh e stencil
s     TypeR e'
t Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
x acc aenv (Array sh e)
a
    stencil2 :: StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> acc aenv (Array sh a1)
-> acc aenv (Array sh b)
-> PreOpenAcc acc aenv (Array sh c)
stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a1)
x Boundary aenv (Array sh b)
y acc aenv (Array sh a1)
a acc aenv (Array sh b)
b  = StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> acc aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> acc aenv (Array sh b)
-> PreOpenAcc acc aenv (Array sh c)
forall sh a1 stencil1 b stencil2 c aenv (acc :: * -> * -> *).
StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> acc aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> acc aenv (Array sh b)
-> PreOpenAcc acc aenv (Array sh c)
Stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a1)
x acc aenv (Array sh a1)
a Boundary aenv (Array sh b)
y acc aenv (Array sh b)
b

    -- Conversions for closed scalar functions and expressions. This just
    -- applies scalar simplifications.
    --
    cvtF :: HasCallStack => Fun aenv' t -> Fun aenv' t
    cvtF :: forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF = Fun aenv' t -> Fun aenv' t
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
simplifyFun

    cvtE :: HasCallStack => Exp aenv' t -> Exp aenv' t
    cvtE :: forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
cvtE = Exp aenv' t -> Exp aenv' t
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
simplifyExp

    cvtB :: HasCallStack => Boundary aenv' t -> Boundary aenv' t
    cvtB :: forall aenv' t.
HasCallStack =>
Boundary aenv' t -> Boundary aenv' t
cvtB Boundary aenv' t
Clamp        = Boundary aenv' t
forall aenv t. Boundary aenv t
Clamp
    cvtB Boundary aenv' t
Mirror       = Boundary aenv' t
forall aenv t. Boundary aenv t
Mirror
    cvtB Boundary aenv' t
Wrap         = Boundary aenv' t
forall aenv t. Boundary aenv t
Wrap
    cvtB (Constant e
c) = e -> Boundary aenv' (Array sh e)
forall e aenv sh. e -> Boundary aenv (Array sh e)
Constant e
c
    cvtB (Function Fun aenv' (sh -> e)
f) = Fun aenv' (sh -> e) -> Boundary aenv' (Array sh e)
forall aenv sh e. Fun aenv (sh -> e) -> Boundary aenv (Array sh e)
Function (Fun aenv' (sh -> e) -> Fun aenv' (sh -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
cvtF Fun aenv' (sh -> e)
f)

    -- Helpers to embed and fuse delayed terms
    --
    into :: (HasCallStack, Sink f)
         => (f env' a -> b)
         -> f env a
         -> Extend ArrayR OpenAcc env env'
         -> b
    into :: forall (f :: * -> * -> *) env' a b env.
(HasCallStack, Sink f) =>
(f env' a -> b) -> f env a -> Extend ArrayR OpenAcc env env' -> b
into f env' a -> b
op f env a
a Extend ArrayR OpenAcc env env'
env = f env' a -> b
op (Extend ArrayR OpenAcc env env' -> f env a -> f env' a
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f env a
a)

    into2 :: (HasCallStack, Sink f1, Sink f2)
          => (f1 env' a -> f2 env' b -> c)
          -> f1 env a
          -> f2 env b
          -> Extend ArrayR OpenAcc env env'
          -> c
    into2 :: forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env.
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> f2 env' b -> c)
-> f1 env a -> f2 env b -> Extend ArrayR OpenAcc env env' -> c
into2 f1 env' a -> f2 env' b -> c
op f1 env a
a f2 env b
b Extend ArrayR OpenAcc env env'
env = f1 env' a -> f2 env' b -> c
op (Extend ArrayR OpenAcc env env' -> f1 env a -> f1 env' a
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f1 env a
a) (Extend ArrayR OpenAcc env env' -> f2 env b -> f2 env' b
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f2 env b
b)

    into2M :: (HasCallStack, Sink f1, Sink f2)
           => (f1 env' a -> Maybe (f2 env' b) -> c)
           -> f1 env a
           -> Maybe (f2 env b)
           -> Extend ArrayR acc env env'
           -> c
    into2M :: forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env
       (acc :: * -> * -> *).
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> Maybe (f2 env' b) -> c)
-> f1 env a -> Maybe (f2 env b) -> Extend ArrayR acc env env' -> c
into2M f1 env' a -> Maybe (f2 env' b) -> c
op f1 env a
a Maybe (f2 env b)
b Extend ArrayR acc env env'
env = f1 env' a -> Maybe (f2 env' b) -> c
op (Extend ArrayR acc env env' -> f1 env a -> f1 env' a
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR acc env env'
env f1 env a
a) (Extend ArrayR acc env env' -> f2 env b -> f2 env' b
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR acc env env'
env (f2 env b -> f2 env' b) -> Maybe (f2 env b) -> Maybe (f2 env' b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (f2 env b)
b)

    into3 :: (HasCallStack, Sink f1, Sink f2, Sink f3)
          => (f1 env' a -> f2 env' b -> f3 env' c -> d)
          -> f1 env a
          -> f2 env b
          -> f3 env c
          -> Extend ArrayR OpenAcc env env'
          -> d
    into3 :: forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) (f3 :: * -> * -> *)
       env' a b c d env.
(HasCallStack, Sink f1, Sink f2, Sink f3) =>
(f1 env' a -> f2 env' b -> f3 env' c -> d)
-> f1 env a
-> f2 env b
-> f3 env c
-> Extend ArrayR OpenAcc env env'
-> d
into3 f1 env' a -> f2 env' b -> f3 env' c -> d
op f1 env a
a f2 env b
b f3 env c
c Extend ArrayR OpenAcc env env'
env = f1 env' a -> f2 env' b -> f3 env' c -> d
op (Extend ArrayR OpenAcc env env' -> f1 env a -> f1 env' a
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f1 env a
a) (Extend ArrayR OpenAcc env env' -> f2 env b -> f2 env' b
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f2 env b
b) (Extend ArrayR OpenAcc env env' -> f3 env c -> f3 env' c
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f3 env c
c)

    -- Operations which can be fused into consumers. Move all of the local
    -- bindings out of the way so that the fusible function operates
    -- directly on the delayed representation. See also: [Representing
    -- delayed arrays]
    --
    fuse :: HasCallStack
         => (forall r aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation D aenv' bs)
         ->       OpenAcc aenv as
         -> Embed OpenAcc aenv bs
    fuse :: forall aenv as bs.
HasCallStack =>
(forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as -> Cunctation D aenv' bs)
-> OpenAcc aenv as -> Embed OpenAcc aenv bs
fuse forall r aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Cunctation D aenv' bs
op (OpenAcc aenv as -> Embed OpenAcc aenv as
EmbedAcc OpenAcc
embedAcc -> Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' as
cc) = Extend ArrayR OpenAcc aenv aenv'
-> Cunctation D aenv' bs -> Embed OpenAcc aenv bs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Cunctation D aenv' bs
forall r aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Cunctation D aenv' bs
op Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' as
cc)

    fuse2 :: HasCallStack
          => (forall r s aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation s aenv' bs -> Cunctation D aenv' cs)
          ->       OpenAcc aenv as
          ->       OpenAcc aenv bs
          -> Embed OpenAcc aenv cs
    fuse2 :: forall aenv as bs cs.
HasCallStack =>
(forall r s aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as
 -> Cunctation s aenv' bs
 -> Cunctation D aenv' cs)
-> OpenAcc aenv as -> OpenAcc aenv bs -> Embed OpenAcc aenv cs
fuse2 forall r s aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as
-> Cunctation s aenv' bs
-> Cunctation D aenv' cs
op OpenAcc aenv as
a1 OpenAcc aenv bs
a0
      | Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' as
cc1  <- OpenAcc aenv as -> Embed OpenAcc aenv as
EmbedAcc OpenAcc
embedAcc OpenAcc aenv as
a1
      , Embed Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' bs
cc0  <- OpenAcc aenv' bs -> Embed OpenAcc aenv' bs
EmbedAcc OpenAcc
embedAcc (Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv bs -> OpenAcc aenv' bs
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv aenv'
env1 OpenAcc aenv bs
a0)
      , Extend ArrayR OpenAcc aenv aenv'
env             <- Extend ArrayR OpenAcc aenv aenv'
env1 Extend ArrayR OpenAcc aenv aenv'
-> Extend ArrayR OpenAcc aenv' aenv'
-> Extend ArrayR OpenAcc aenv aenv'
forall (s :: * -> *) (acc :: * -> * -> *) env env' env''.
Extend s acc env env'
-> Extend s acc env' env'' -> Extend s acc env env''
`append` Extend ArrayR OpenAcc aenv' aenv'
env0
      = Extend ArrayR OpenAcc aenv aenv'
-> Cunctation D aenv' cs -> Embed OpenAcc aenv cs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as
-> Cunctation r aenv' bs
-> Cunctation D aenv' cs
forall r s aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as
-> Cunctation s aenv' bs
-> Cunctation D aenv' cs
op Extend ArrayR OpenAcc aenv aenv'
env (Extend ArrayR OpenAcc aenv' aenv'
-> Cunctation r aenv' as -> Cunctation r aenv' as
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' as
cc1) Cunctation r aenv' bs
cc0)

    -- Consumer operations which will be evaluated.
    --
    -- NOTE: [Fusion and the lowest common use site]
    --
    -- The AST given to us by sharing recovery will place let bindings at
    -- the lowest common use site for that shared term. For example:
    --
    --   fold f z (let a0 = ..
    --                 a1 = ..
    --              in zipWith g a0 a1)
    --
    -- In order to enable producer/consumer fusion for the above example,
    -- it is necessary to float the let bindings above the `fold`
    -- operation; SEE: [Sharing vs. Fusion] for more information.
    --
    -- Furthermore, we used to maintain an invariant that all (manifest)
    -- arguments were supplied as array variables, for example:
    --
    --   fold1 f (let a0 = ..               let a0 = ..
    --             in stencil g a0)   ==>       a1 = stencil g a0
    --                                          a2 = fold1 f a1
    --
    -- However, if the argument term will be evaluated (i.e. can not be
    -- fused into the producer) then it is better that we do _not_ float
    -- those terms, and instead leave them under the consumer. This helps
    -- to syntactically constrain the "liveness" of terms: if the argument
    -- to an operation is not an array variable, we can see directly that
    -- this will be the last use-site of that array. In particular, this is
    -- useful for the 'permute' operation to know when it can in-place
    -- update the array of default values.
    --
    embed :: HasCallStack
          => ArraysR bs
          -> (forall aenv'. Extend ArrayR OpenAcc aenv aenv' -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs)
          ->       OpenAcc aenv as
          -> Embed OpenAcc aenv bs
    embed :: forall bs aenv as.
HasCallStack =>
ArraysR bs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs)
-> OpenAcc aenv as
-> Embed OpenAcc aenv bs
embed ArraysR bs
reprBs forall aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs
op (OpenAcc aenv as -> Embed OpenAcc aenv as
EmbedAcc OpenAcc
embedAcc -> Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' as
cc)
      | Done{} <- Cunctation r aenv' as
cc
      , DeclareVars LeftHandSide ArrayR bs aenv env'
lhs aenv :> env'
_ forall env''. (env' :> env'') -> Vars ArrayR env'' bs
value <- ArraysR bs -> DeclareVars ArrayR bs aenv
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars ArraysR bs
reprBs
      = Extend ArrayR OpenAcc aenv env'
-> Cunctation M env' bs -> Embed OpenAcc aenv bs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv
-> LeftHandSide ArrayR bs aenv env'
-> OpenAcc aenv bs
-> Extend ArrayR OpenAcc aenv env'
forall (s :: * -> *) (f :: * -> * -> *) env env'1 t env'.
Extend s f env env'1
-> LeftHandSide s t env'1 env' -> f env'1 t -> Extend s f env env'
PushEnv Extend ArrayR OpenAcc aenv aenv
forall (s :: * -> *) (f :: * -> * -> *) env. Extend s f env env
BaseEnv LeftHandSide ArrayR bs aenv env'
lhs (OpenAcc aenv bs -> Extend ArrayR OpenAcc aenv env')
-> OpenAcc aenv bs -> Extend ArrayR OpenAcc aenv env'
forall a b. (a -> b) -> a -> b
$ PreOpenAcc OpenAcc aenv bs -> OpenAcc aenv bs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Extend ArrayR OpenAcc aenv aenv
-> OpenAcc aenv as -> PreOpenAcc OpenAcc aenv bs
forall aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs
op Extend ArrayR OpenAcc aenv aenv
forall (s :: * -> *) (f :: * -> * -> *) env. Extend s f env env
BaseEnv (Embed OpenAcc aenv as -> OpenAcc aenv as
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Embed OpenAcc aenv as
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' as
cc)))) (Cunctation M env' bs -> Embed OpenAcc aenv bs)
-> Cunctation M env' bs -> Embed OpenAcc aenv bs
forall a b. (a -> b) -> a -> b
$ ArrayVars env' bs -> Cunctation M env' bs
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done (ArrayVars env' bs -> Cunctation M env' bs)
-> ArrayVars env' bs -> Cunctation M env' bs
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ArrayVars env' bs
forall env''. (env' :> env'') -> Vars ArrayR env'' bs
value env' :> env'
forall env. env :> env
weakenId
      | Bool
otherwise
      -- Next line is duplicated for both branches, as the type variable for the environment is instantiated differently
      , DeclareVars LeftHandSide ArrayR bs aenv' env'
lhs aenv' :> env'
_ forall env''. (env' :> env'') -> Vars ArrayR env'' bs
value <- ArraysR bs -> DeclareVars ArrayR bs aenv'
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars ArraysR bs
reprBs
      = Extend ArrayR OpenAcc aenv env'
-> Cunctation M env' bs -> Embed OpenAcc aenv bs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
-> LeftHandSide ArrayR bs aenv' env'
-> OpenAcc aenv' bs
-> Extend ArrayR OpenAcc aenv env'
forall (s :: * -> *) (f :: * -> * -> *) env env'1 t env'.
Extend s f env env'1
-> LeftHandSide s t env'1 env' -> f env'1 t -> Extend s f env env'
PushEnv Extend ArrayR OpenAcc aenv aenv'
env     LeftHandSide ArrayR bs aenv' env'
lhs (OpenAcc aenv' bs -> Extend ArrayR OpenAcc aenv env')
-> OpenAcc aenv' bs -> Extend ArrayR OpenAcc aenv env'
forall a b. (a -> b) -> a -> b
$ PreOpenAcc OpenAcc aenv' bs -> OpenAcc aenv' bs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs
forall aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs
op Extend ArrayR OpenAcc aenv aenv'
env     (PreOpenAcc OpenAcc aenv' as -> OpenAcc aenv' as
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Cunctation r aenv' as -> PreOpenAcc OpenAcc aenv' as
forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute Cunctation r aenv' as
cc))))      (Cunctation M env' bs -> Embed OpenAcc aenv bs)
-> Cunctation M env' bs -> Embed OpenAcc aenv bs
forall a b. (a -> b) -> a -> b
$ ArrayVars env' bs -> Cunctation M env' bs
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done (ArrayVars env' bs -> Cunctation M env' bs)
-> ArrayVars env' bs -> Cunctation M env' bs
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ArrayVars env' bs
forall env''. (env' :> env'') -> Vars ArrayR env'' bs
value env' :> env'
forall env. env :> env
weakenId

    embed2 :: HasCallStack
           => ArraysR cs
           -> (forall aenv'. Extend ArrayR OpenAcc aenv aenv' -> OpenAcc aenv' as -> OpenAcc aenv' bs -> PreOpenAcc OpenAcc aenv' cs)
           ->       OpenAcc aenv as
           ->       OpenAcc aenv bs
           -> Embed OpenAcc aenv cs
    embed2 :: forall cs aenv as bs.
HasCallStack =>
ArraysR cs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as
    -> OpenAcc aenv' bs
    -> PreOpenAcc OpenAcc aenv' cs)
-> OpenAcc aenv as
-> OpenAcc aenv bs
-> Embed OpenAcc aenv cs
embed2 ArraysR cs
reprCs forall aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as
-> OpenAcc aenv' bs
-> PreOpenAcc OpenAcc aenv' cs
op (OpenAcc aenv as -> Embed OpenAcc aenv as
EmbedAcc OpenAcc
embedAcc -> Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' as
cc1) OpenAcc aenv bs
a0
      | Done{}          <- Cunctation r aenv' as
cc1
      , OpenAcc aenv as
a1              <- Embed OpenAcc aenv as -> OpenAcc aenv as
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Embed OpenAcc aenv as
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' as
cc1)
      = ArraysR cs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' bs -> PreOpenAcc OpenAcc aenv' cs)
-> OpenAcc aenv bs
-> Embed OpenAcc aenv cs
forall bs aenv as.
HasCallStack =>
ArraysR bs
-> (forall aenv'.
    Extend ArrayR OpenAcc aenv aenv'
    -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs)
-> OpenAcc aenv as
-> Embed OpenAcc aenv bs
embed ArraysR cs
reprCs (\Extend ArrayR OpenAcc aenv aenv'
env0 -> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as
-> OpenAcc aenv' bs
-> PreOpenAcc OpenAcc aenv' cs
forall aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as
-> OpenAcc aenv' bs
-> PreOpenAcc OpenAcc aenv' cs
op Extend ArrayR OpenAcc aenv aenv'
env0 (Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv as -> OpenAcc aenv' as
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv aenv'
env0 OpenAcc aenv as
a1)) OpenAcc aenv bs
a0
      --
      | Embed Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' bs
cc0  <- OpenAcc aenv' bs -> Embed OpenAcc aenv' bs
EmbedAcc OpenAcc
embedAcc (Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv bs -> OpenAcc aenv' bs
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv aenv'
env1 OpenAcc aenv bs
a0)
      , Extend ArrayR OpenAcc aenv aenv'
env             <- Extend ArrayR OpenAcc aenv aenv'
env1 Extend ArrayR OpenAcc aenv aenv'
-> Extend ArrayR OpenAcc aenv' aenv'
-> Extend ArrayR OpenAcc aenv aenv'
forall (s :: * -> *) (acc :: * -> * -> *) env env' env''.
Extend s acc env env'
-> Extend s acc env' env'' -> Extend s acc env env''
`append` Extend ArrayR OpenAcc aenv' aenv'
env0
      = case Cunctation r aenv' bs
cc0 of
          Done{}
            | DeclareVars LeftHandSide ArrayR cs aenv' env'
lhs aenv' :> env'
_ forall env''. (env' :> env'') -> Vars ArrayR env'' cs
value <- ArraysR cs -> DeclareVars ArrayR cs aenv'
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars ArraysR cs
reprCs
              -> Extend ArrayR OpenAcc aenv env'
-> Cunctation M env' cs -> Embed OpenAcc aenv cs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
-> LeftHandSide ArrayR cs aenv' env'
-> OpenAcc aenv' cs
-> Extend ArrayR OpenAcc aenv env'
forall (s :: * -> *) (f :: * -> * -> *) env env'1 t env'.
Extend s f env env'1
-> LeftHandSide s t env'1 env' -> f env'1 t -> Extend s f env env'
PushEnv Extend ArrayR OpenAcc aenv aenv'
env1 LeftHandSide ArrayR cs aenv' env'
lhs (OpenAcc aenv' cs -> Extend ArrayR OpenAcc aenv env')
-> OpenAcc aenv' cs -> Extend ArrayR OpenAcc aenv env'
forall a b. (a -> b) -> a -> b
$ PreOpenAcc OpenAcc aenv' cs -> OpenAcc aenv' cs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as
-> OpenAcc aenv' bs
-> PreOpenAcc OpenAcc aenv' cs
forall aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as
-> OpenAcc aenv' bs
-> PreOpenAcc OpenAcc aenv' cs
op Extend ArrayR OpenAcc aenv aenv'
env1 (PreOpenAcc OpenAcc aenv' as -> OpenAcc aenv' as
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Cunctation r aenv' as -> PreOpenAcc OpenAcc aenv' as
forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute Cunctation r aenv' as
cc1)) (Embed OpenAcc aenv' bs -> OpenAcc aenv' bs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Extend ArrayR OpenAcc aenv' aenv'
-> Cunctation r aenv' bs -> Embed OpenAcc aenv' bs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' bs
cc0))))
                       (Cunctation M env' cs -> Embed OpenAcc aenv cs)
-> Cunctation M env' cs -> Embed OpenAcc aenv cs
forall a b. (a -> b) -> a -> b
$ ArrayVars env' cs -> Cunctation M env' cs
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done
                       (ArrayVars env' cs -> Cunctation M env' cs)
-> ArrayVars env' cs -> Cunctation M env' cs
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ArrayVars env' cs
forall env''. (env' :> env'') -> Vars ArrayR env'' cs
value env' :> env'
forall env. env :> env
weakenId
          Cunctation r aenv' bs
_
            -- Next line is duplicated for both branches, as the type
            -- variable for the environment is instantiated differently
            | DeclareVars LeftHandSide ArrayR cs aenv' env'
lhs aenv' :> env'
_ forall env''. (env' :> env'') -> Vars ArrayR env'' cs
value <- ArraysR cs -> DeclareVars ArrayR cs aenv'
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars ArraysR cs
reprCs
              -> Extend ArrayR OpenAcc aenv env'
-> Cunctation M env' cs -> Embed OpenAcc aenv cs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
-> LeftHandSide ArrayR cs aenv' env'
-> OpenAcc aenv' cs
-> Extend ArrayR OpenAcc aenv env'
forall (s :: * -> *) (f :: * -> * -> *) env env'1 t env'.
Extend s f env env'1
-> LeftHandSide s t env'1 env' -> f env'1 t -> Extend s f env env'
PushEnv Extend ArrayR OpenAcc aenv aenv'
env  LeftHandSide ArrayR cs aenv' env'
lhs (OpenAcc aenv' cs -> Extend ArrayR OpenAcc aenv env')
-> OpenAcc aenv' cs -> Extend ArrayR OpenAcc aenv env'
forall a b. (a -> b) -> a -> b
$ PreOpenAcc OpenAcc aenv' cs -> OpenAcc aenv' cs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as
-> OpenAcc aenv' bs
-> PreOpenAcc OpenAcc aenv' cs
forall aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' as
-> OpenAcc aenv' bs
-> PreOpenAcc OpenAcc aenv' cs
op Extend ArrayR OpenAcc aenv aenv'
env  (PreOpenAcc OpenAcc aenv' as -> OpenAcc aenv' as
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Cunctation r aenv' as -> PreOpenAcc OpenAcc aenv' as
forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute (Extend ArrayR OpenAcc aenv' aenv'
-> Cunctation r aenv' as -> Cunctation r aenv' as
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' as
cc1))) (PreOpenAcc OpenAcc aenv' bs -> OpenAcc aenv' bs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Cunctation r aenv' bs -> PreOpenAcc OpenAcc aenv' bs
forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute Cunctation r aenv' bs
cc0))))
                       (Cunctation M env' cs -> Embed OpenAcc aenv cs)
-> Cunctation M env' cs -> Embed OpenAcc aenv cs
forall a b. (a -> b) -> a -> b
$ ArrayVars env' cs -> Cunctation M env' cs
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done
                       (ArrayVars env' cs -> Cunctation M env' cs)
-> ArrayVars env' cs -> Cunctation M env' cs
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ArrayVars env' cs
forall env''. (env' :> env'') -> Vars ArrayR env'' cs
value env' :> env'
forall env. env :> env
weakenId

    -- trav1 :: (Arrays as, Arrays bs)
    --       => (forall aenv'. Embed acc aenv' as -> Embed acc aenv' as)
    --       -> (forall aenv'. Extend ArrayR acc aenv aenv' -> acc aenv' as -> PreOpenAcc acc aenv' bs)
    --       ->       acc aenv as
    --       -> Embed acc aenv bs
    -- trav1 f op (f . embedAcc -> Embed env cc)
    --   = Embed (env `pushArrayEnv` inject (op env (inject (compute cc)))) doneZeroIdx

    -- trav2 :: (Arrays as, Arrays bs, Arrays cs)
    --       => (forall aenv'. Embed acc aenv' as -> Embed acc aenv' as)
    --       -> (forall aenv'. Embed acc aenv' bs -> Embed acc aenv' bs)
    --       -> (forall aenv'. Extend ArrayR acc aenv aenv' -> acc aenv' as -> acc aenv' bs -> PreOpenAcc acc aenv' cs)
    --       ->       acc aenv as
    --       ->       acc aenv bs
    --       -> Embed acc aenv cs
    -- trav2 f1 f0 op (f1 . embedAcc -> Embed env1 cc1) (f0 . embedAcc . sinkA env1 -> Embed env0 cc0)
    --   | env     <- env1 `append` env0
    --   , acc1    <- inject . compute $ sinkA env0 cc1
    --   , acc0    <- inject . compute $ cc0
    --   = Embed (env `pushArrayEnv` inject (op env acc1 acc0)) doneZeroIdx

    -- force :: Arrays as => Embed acc aenv' as -> Embed acc aenv' as
    -- force (Embed env cc)
    --   | Done{} <- cc = Embed env                                  cc
    --   | otherwise    = Embed (env `pushArrayEnv` inject (compute cc)) doneZeroIdx

    -- -- Move additional bindings for producers outside of the sequence, so that
    -- -- producers may fuse with their arguments resulting in actual sequencing
    -- collectD :: PreOpenSeq acc aenv () arrs
    --          -> Embed acc aenv arrs
    -- collectD (embedSeq embedAcc -> ExtendSeq env s')
    --   = Embed (env `pushArrayEnv` inject (Collect s')) doneZeroIdx


{--
-- Move additional bindings for producer outside of sequence, so
-- that producers may fuse with their arguments, resulting in
-- actual sequencing.
embedSeq :: forall acc aenv arrs. Kit acc
         => EmbedAcc acc
         -> PreOpenSeq acc aenv () arrs
         -> ExtendSeq       acc aenv () arrs
embedSeq embedAcc s
  = travS s BaseEnv
  where
    travS :: forall senv aenv' arrs'.
             PreOpenSeq acc aenv senv arrs'
          -> Extend acc aenv aenv'
          -> ExtendSeq acc aenv senv arrs'
    travS s env =
      case s of
        Producer p s
          | ExtendSeq env' s' <- travS s env
          , ExtendProducer env'' p' <- travP p env'
          -> ExtendSeq (env' `append` env'') (Producer p' (sinkSeq env'' s'))
        Consumer c
          | c' <- travC c env
          -> ExtendSeq env (Consumer c')
        Reify ix
          -> ExtendSeq env (Reify ix)

    travP :: forall arrs' aenv' senv.
             Producer acc aenv senv arrs'
          -> Extend acc aenv aenv'
          -> ExtendProducer acc aenv' senv arrs'
    travP (ToSeq slix sh a) env
      | Embed env' cc <- embedAcc (sink env a)
      = ExtendProducer env' (ToSeq slix sh (inject (compute cc)))
    travP (StreamIn arrs) _          = ExtendProducer BaseEnv (StreamIn arrs)
    travP (MapSeq f x) env           = ExtendProducer BaseEnv (MapSeq (cvtAF (sink env f)) x)
    travP (ChunkedMapSeq f x) env    = ExtendProducer BaseEnv (ChunkedMapSeq (cvtAF (sink env f)) x)
    travP (ZipWithSeq f x y) env     = ExtendProducer BaseEnv (ZipWithSeq (cvtAF (sink env f)) x y)
    travP (ScanSeq f e x) env        = ExtendProducer BaseEnv (ScanSeq (cvtF (sink env f)) (cvtE (sink env e)) x)

    travC :: forall arrs' aenv' senv.
             Consumer acc aenv senv arrs'
          -> Extend acc aenv aenv'
          -> Consumer acc aenv' senv arrs'
    travC (FoldSeq f e x) env = FoldSeq (cvtF (sink env f)) (cvtE (sink env e)) x
    travC (FoldSeqFlatten f a x) env = FoldSeqFlatten (cvtAF (sink env f)) (cvtA (sink env a)) x
    travC (Stuple t) env = Stuple (cvtCT t)
      where
        cvtCT :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv' senv) t
        cvtCT NilAtup        = NilAtup
        cvtCT (SnocAtup t c) = SnocAtup (cvtCT t) (travC c env)

    cvtE :: Elt t => Exp aenv' t -> Exp aenv' t
    cvtE = simplifyExp

    cvtF :: Fun aenv' t -> Fun aenv' t
    cvtF = simplifyFun

    cvtA :: Arrays a => acc aenv' a -> acc aenv' a
    cvtA = computeAcc . embedAcc

    cvtAF :: PreOpenAfun acc aenv' f -> PreOpenAfun acc aenv' f
    cvtAF (Alam  f) = Alam  (cvtAF f)
    cvtAF (Abody a) = Abody (cvtA a)


-- A sequence with additional bindings
data ExtendSeq acc aenv senv arrs where
  ExtendSeq :: forall acc aenv aenv' senv arrs.
                Extend acc aenv aenv'
             -> PreOpenSeq acc aenv' senv arrs
             -> ExtendSeq acc aenv senv arrs

-- A producer with additional bindings
data ExtendProducer acc aenv senv arrs where
  ExtendProducer :: forall acc aenv aenv' senv arrs.
                    Extend acc aenv aenv'
                 -> Producer acc aenv' senv arrs
                 -> ExtendProducer acc aenv senv arrs
--}


-- Internal representation
-- =======================

-- NOTE: [Representing delayed arrays]
--
-- During the fusion transformation we represent terms as a pair consisting of
-- a collection of supplementary environment bindings and a description of how
-- to construct the array.
--
-- It is critical to separate these two. To create a real AST node we need both
-- the environment and array term, but analysis of how to fuse terms requires
-- only the array description. If the additional bindings are bundled as part of
-- the representation, the existentially quantified extended environment type
-- will be untouchable. This is problematic because the terms of the two arrays
-- are defined with respect to this existentially quantified type, and there is
-- no way to directly combine these two environments:
--
--   append :: Extend env env1 -> Extend env env2 -> Extend env ???
--
-- And hence, no way to combine the terms of the delayed representation.
--
-- The only way to bring terms into the same scope is to operate via the
-- manifest terms. This entails a great deal of conversion between delayed and
-- AST terms, but is certainly possible.
--
-- However, because of the limited scope into which this existential type is
-- available, we ultimately perform this process many times. In fact, complexity
-- of the fusion algorithm for an AST of N terms becomes O(r^n), where r is the
-- number of different rules we have for combining terms.
--
data Embed acc aenv a where
  Embed :: Extend ArrayR acc aenv aenv'
        -> Cunctation r           aenv' a
        -> Embed         acc aenv       a

instance HasArraysR acc => HasArraysR (Embed acc) where
  arraysR :: forall aenv a. Embed acc aenv a -> ArraysR a
arraysR (Embed Extend ArrayR acc aenv aenv'
_ Cunctation r aenv' a
c) = Cunctation r aenv' a -> ArraysR a
forall aenv a. Cunctation r aenv a -> ArraysR a
forall (f :: * -> * -> *) aenv a.
HasArraysR f =>
f aenv a -> ArraysR a
arraysR Cunctation r aenv' a
c


-- Delayed arrays are represented as functions
data D

-- Manifest arrays are real arrays backed by unboxed vectors
data M

-- Cunctation (n): the action or an instance of delaying; a tardy action.
--
-- This describes the ways in which the fusion transformation represents
-- intermediate arrays. The fusion process operates by recasting producer array
-- computations in terms of a set of scalar functions used to construct an
-- element at each index, and fusing successive producers by combining these
-- scalar functions.
--
data Cunctation r aenv a where

  -- The base case is just a real (manifest) array term. No fusion happens here.
  -- Note that the array is referenced by an index into the extended
  -- environment, ensuring that the array is manifest and making the term
  -- non-recursive in 'acc'.
  --
  Done  :: ArrayVars    aenv arrs
        -> Cunctation M aenv arrs

  -- We can represent an array by its shape and a function to compute an element
  -- at each index.
  --
  Yield :: ArrayR (Array sh e)
        -> Exp          aenv sh
        -> Fun          aenv (sh -> e)
        -> Cunctation D aenv (Array sh e)

  -- A more restrictive form than 'Yield' may afford greater opportunities for
  -- optimisation by a backend. This more structured form applies an index and
  -- value transform to an input array. Note that the transform is applied to an
  -- array stored as an environment index, so that the term is non-recursive and
  -- it is always possible to embed into a collective operation.
  --
  Step  :: ArrayR (Array sh' b)
        -> Exp          aenv sh'
        -> Fun          aenv (sh' -> sh)
        -> Fun          aenv (a   -> b)
        -> ArrayVar     aenv (Array sh  a)
        -> Cunctation D aenv (Array sh' b)

instance HasArraysR (Cunctation r) where
  arraysR :: forall aenv a. Cunctation r aenv a -> ArraysR a
arraysR (Done ArrayVars aenv a
v)          = ArrayVars aenv a -> TupR ArrayR a
forall (s :: * -> *) env t. Vars s env t -> TupR s t
varsType ArrayVars aenv a
v
  arraysR (Yield ArrayR (Array sh e)
aR Exp aenv sh
_ Fun aenv (sh -> e)
_)    = ArrayR a -> TupR ArrayR a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR a
ArrayR (Array sh e)
aR
  arraysR (Step ArrayR (Array sh' b)
aR Exp aenv sh'
_ Fun aenv (sh' -> sh)
_ Fun aenv (a -> b)
_ ArrayVar aenv (Array sh a)
_) = ArrayR a -> TupR ArrayR a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR a
ArrayR (Array sh' b)
aR

instance Sink (Cunctation r) where
  weaken :: forall env env'.
(env :> env')
-> forall t. Cunctation r env t -> Cunctation r env' t
weaken env :> env'
k = \case
    Done ArrayVars env t
v            -> ArrayVars env' t -> Cunctation M env' t
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done ((env :> env') -> ArrayVars env t -> ArrayVars env' t
forall env env' (s :: * -> *) t.
(env :> env') -> Vars s env t -> Vars s env' t
weakenVars env :> env'
k ArrayVars env t
v)
    Step ArrayR (Array sh' b)
aR Exp env sh'
sh Fun env (sh' -> sh)
p Fun env (a -> b)
f ArrayVar env (Array sh a)
v  -> ArrayR (Array sh' b)
-> Exp env' sh'
-> Fun env' (sh' -> sh)
-> Fun env' (a -> b)
-> ArrayVar env' (Array sh a)
-> Cunctation D env' (Array sh' b)
forall env' e aenv sh a.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> sh)
-> Fun aenv (a -> e)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array env' e)
Step  ArrayR (Array sh' b)
aR ((env :> env') -> forall t. OpenExp () env t -> OpenExp () env' t
forall env env'.
(env :> env') -> forall t. OpenExp () env t -> OpenExp () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken env :> env'
k Exp env sh'
sh) ((env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall env env'.
(env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken env :> env'
k Fun env (sh' -> sh)
p) ((env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall env env'.
(env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken env :> env'
k Fun env (a -> b)
f) ((env :> env') -> forall t. Var ArrayR env t -> Var ArrayR env' t
forall env env'.
(env :> env') -> forall t. Var ArrayR env t -> Var ArrayR env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken env :> env'
k ArrayVar env (Array sh a)
v)
    Yield ArrayR (Array sh e)
aR Exp env sh
sh Fun env (sh -> e)
f     -> ArrayR (Array sh e)
-> Exp env' sh
-> Fun env' (sh -> e)
-> Cunctation D env' (Array sh e)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield ArrayR (Array sh e)
aR ((env :> env') -> forall t. OpenExp () env t -> OpenExp () env' t
forall env env'.
(env :> env') -> forall t. OpenExp () env t -> OpenExp () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken env :> env'
k Exp env sh
sh) ((env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall env env'.
(env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken env :> env'
k Fun env (sh -> e)
f)

simplify
    :: HasCallStack
    => Cunctation r aenv a
    -> Either (Cunctation M aenv a) (Cunctation D aenv a)
simplify :: forall r aenv a.
HasCallStack =>
Cunctation r aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
simplify = \case
  Done ArrayVars aenv a
v
    -> Cunctation M aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. a -> Either a b
Left  (Cunctation M aenv a
 -> Either (Cunctation M aenv a) (Cunctation D aenv a))
-> Cunctation M aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. (a -> b) -> a -> b
$ ArrayVars aenv a -> Cunctation M aenv a
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done ArrayVars aenv a
v
  Yield ArrayR (Array sh e)
aR (Exp aenv sh -> Exp aenv sh
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
simplifyExp -> Exp aenv sh
sh) (Fun aenv (sh -> e) -> Fun aenv (sh -> e)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
simplifyFun -> Fun aenv (sh -> e)
f)
    -> Cunctation D aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. b -> Either a b
Right (Cunctation D aenv a
 -> Either (Cunctation M aenv a) (Cunctation D aenv a))
-> Cunctation D aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Cunctation D aenv (Array sh e)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield ArrayR (Array sh e)
aR Exp aenv sh
sh Fun aenv (sh -> e)
f
  Step ArrayR (Array sh' b)
aR (Exp aenv sh' -> Exp aenv sh'
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
simplifyExp -> Exp aenv sh'
sh) (Fun aenv (sh' -> sh) -> Fun aenv (sh' -> sh)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
simplifyFun -> Fun aenv (sh' -> sh)
p) (Fun aenv (a -> b) -> Fun aenv (a -> b)
forall aenv' t. HasCallStack => Fun aenv' t -> Fun aenv' t
simplifyFun -> Fun aenv (a -> b)
f) ArrayVar aenv (Array sh a)
v
    | Just sh' :~: sh
Refl <- Exp aenv sh' -> OpenExp () aenv sh -> Maybe (sh' :~: sh)
forall env aenv s t.
OpenExp env aenv s -> OpenExp env aenv t -> Maybe (s :~: t)
matchOpenExp Exp aenv sh'
sh (ArrayVar aenv (Array sh a) -> OpenExp () aenv sh
forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape ArrayVar aenv (Array sh a)
v)
    , Just sh' :~: sh
Refl <- Fun aenv (sh' -> sh) -> Maybe (sh' :~: sh)
forall env aenv a b. OpenFun env aenv (a -> b) -> Maybe (a :~: b)
isIdentity Fun aenv (sh' -> sh)
p
    , Just a :~: b
Refl <- Fun aenv (a -> b) -> Maybe (a :~: b)
forall env aenv a b. OpenFun env aenv (a -> b) -> Maybe (a :~: b)
isIdentity Fun aenv (a -> b)
f
    -> Cunctation M aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. a -> Either a b
Left  (Cunctation M aenv a
 -> Either (Cunctation M aenv a) (Cunctation D aenv a))
-> Cunctation M aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. (a -> b) -> a -> b
$ (ArrayVars aenv a -> Cunctation M aenv a
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done (Var ArrayR aenv a -> ArrayVars aenv a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle Var ArrayR aenv a
ArrayVar aenv (Array sh a)
v))
    | Bool
otherwise
    -> Cunctation D aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. b -> Either a b
Right (Cunctation D aenv a
 -> Either (Cunctation M aenv a) (Cunctation D aenv a))
-> Cunctation D aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a -> b)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array sh' b)
forall env' e aenv sh a.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> sh)
-> Fun aenv (a -> e)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array env' e)
Step ArrayR (Array sh' b)
aR Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a -> b)
f ArrayVar aenv (Array sh a)
v


-- Convert a real AST node into the internal representation
--
done :: HasCallStack => PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done :: forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done PreOpenAcc OpenAcc aenv a
pacc
  | Just ArrayVars aenv a
vars <- ExtractAcc OpenAcc
-> PreOpenAcc OpenAcc aenv a -> Maybe (ArrayVars aenv a)
forall (acc :: * -> * -> *) aenv a.
ExtractAcc acc -> PreOpenAcc acc aenv a -> Maybe (ArrayVars aenv a)
avarsOut OpenAcc env t -> Maybe (PreOpenAcc OpenAcc env t)
ExtractAcc OpenAcc
extractOpenAcc PreOpenAcc OpenAcc aenv a
pacc
  = Extend ArrayR OpenAcc aenv aenv
-> Cunctation M aenv a -> Embed OpenAcc aenv a
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv
forall (s :: * -> *) (f :: * -> * -> *) env. Extend s f env env
BaseEnv (ArrayVars aenv a -> Cunctation M aenv a
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done ArrayVars aenv a
vars)
  | DeclareVars LeftHandSide ArrayR a aenv env'
lhs aenv :> env'
_ forall env''. (env' :> env'') -> Vars ArrayR env'' a
value <- TupR ArrayR a -> DeclareVars ArrayR a aenv
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (PreOpenAcc OpenAcc aenv a -> TupR ArrayR a
forall aenv a. PreOpenAcc OpenAcc aenv a -> ArraysR a
forall (f :: * -> * -> *) aenv a.
HasArraysR f =>
f aenv a -> ArraysR a
arraysR PreOpenAcc OpenAcc aenv a
pacc)
  = Extend ArrayR OpenAcc aenv env'
-> Cunctation M env' a -> Embed OpenAcc aenv a
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv
-> LeftHandSide ArrayR a aenv env'
-> OpenAcc aenv a
-> Extend ArrayR OpenAcc aenv env'
forall (s :: * -> *) (f :: * -> * -> *) env env'1 t env'.
Extend s f env env'1
-> LeftHandSide s t env'1 env' -> f env'1 t -> Extend s f env env'
PushEnv Extend ArrayR OpenAcc aenv aenv
forall (s :: * -> *) (f :: * -> * -> *) env. Extend s f env env
BaseEnv LeftHandSide ArrayR a aenv env'
lhs (OpenAcc aenv a -> Extend ArrayR OpenAcc aenv env')
-> OpenAcc aenv a -> Extend ArrayR OpenAcc aenv env'
forall a b. (a -> b) -> a -> b
$ PreOpenAcc OpenAcc aenv a -> OpenAcc aenv a
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc PreOpenAcc OpenAcc aenv a
pacc) (Cunctation M env' a -> Embed OpenAcc aenv a)
-> Cunctation M env' a -> Embed OpenAcc aenv a
forall a b. (a -> b) -> a -> b
$ ArrayVars env' a -> Cunctation M env' a
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done (ArrayVars env' a -> Cunctation M env' a)
-> ArrayVars env' a -> Cunctation M env' a
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ArrayVars env' a
forall env''. (env' :> env'') -> Vars ArrayR env'' a
value env' :> env'
forall env. env :> env
weakenId

doneZeroIdx :: ArrayR (Array sh e) -> Cunctation M (aenv, Array sh e) (Array sh e)
doneZeroIdx :: forall sh e aenv.
ArrayR (Array sh e) -> Cunctation M (aenv, Array sh e) (Array sh e)
doneZeroIdx ArrayR (Array sh e)
aR = ArrayVars (aenv, Array sh e) (Array sh e)
-> Cunctation M (aenv, Array sh e) (Array sh e)
forall aenv arrs. ArrayVars aenv arrs -> Cunctation M aenv arrs
Done (Var ArrayR (aenv, Array sh e) (Array sh e)
-> ArrayVars (aenv, Array sh e) (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh e)
-> Idx (aenv, Array sh e) (Array sh e)
-> Var ArrayR (aenv, Array sh e) (Array sh e)
forall (s :: * -> *) env t. s t -> Idx env t -> Var s env t
Var ArrayR (Array sh e)
aR Idx (aenv, Array sh e) (Array sh e)
forall envt t env. (envt ~ (env, t)) => Idx envt t
ZeroIdx))


-- Recast a cunctation into a mapping from indices to elements
--
yield :: HasCallStack
      => Cunctation r aenv (Array sh e)
      -> Cunctation D aenv (Array sh e)
yield :: forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
yield Cunctation r aenv (Array sh e)
cc =
  case Cunctation r aenv (Array sh e)
cc of
    Yield{}                        -> Cunctation r aenv (Array sh e)
Cunctation D aenv (Array sh e)
cc
    Step ArrayR (Array sh' b)
tR Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a -> b)
f ArrayVar aenv (Array sh a)
v               -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Cunctation D aenv (Array sh e)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield ArrayR (Array sh e)
ArrayR (Array sh' b)
tR Exp aenv sh
Exp aenv sh'
sh (OpenFun () aenv (a -> e)
Fun aenv (a -> b)
f OpenFun () aenv (a -> e)
-> OpenFun () aenv (sh -> a) -> Fun aenv (sh -> e)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ArrayVar aenv (Array sh a) -> Fun aenv (sh -> a)
forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray ArrayVar aenv (Array sh a)
v Fun aenv (sh -> a)
-> OpenFun () aenv (sh -> sh) -> OpenFun () aenv (sh -> a)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` OpenFun () aenv (sh -> sh)
Fun aenv (sh' -> sh)
p)
    Done (TupRsingle v :: Var ArrayR aenv (Array sh e)
v@(Var ArrayR (Array sh e)
tR Idx aenv (Array sh e)
_)) -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Cunctation D aenv (Array sh e)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield ArrayR (Array sh e)
tR (Var ArrayR aenv (Array sh e) -> Exp aenv sh
forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape Var ArrayR aenv (Array sh e)
v) (Var ArrayR aenv (Array sh e) -> Fun aenv (sh -> e)
forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray Var ArrayR aenv (Array sh e)
v)

-- Recast a cunctation into a delayed representation
--
delaying
    :: HasCallStack
    => Cunctation r aenv (Array sh e)
    -> Cunctation D aenv (Array sh e)
delaying :: forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
delaying Cunctation r aenv (Array sh e)
cc =
  case Cunctation r aenv (Array sh e)
cc of
    Yield{} -> Cunctation r aenv (Array sh e)
Cunctation D aenv (Array sh e)
cc
    Step{}  -> Cunctation r aenv (Array sh e)
Cunctation D aenv (Array sh e)
cc
    Done ArrayVars aenv (Array sh e)
u
      | TupRsingle Var ArrayR aenv (Array sh e)
v  <- ArrayVars aenv (Array sh e)
u
      , Var ArrayR (Array sh e)
aR Idx aenv (Array sh e)
_      <- Var ArrayR aenv (Array sh e)
v
      , ArrayR ShapeR sh
shR TypeR e
tR <- ArrayR (Array sh e)
aR
     -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> sh)
-> Fun aenv (e -> e)
-> Var ArrayR aenv (Array sh e)
-> Cunctation D aenv (Array sh e)
forall env' e aenv sh a.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> sh)
-> Fun aenv (a -> e)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array env' e)
Step ArrayR (Array sh e)
aR (Var ArrayR aenv (Array sh e) -> Exp aenv sh
forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape Var ArrayR aenv (Array sh e)
v) (TypeR sh -> Fun aenv (sh -> sh)
forall a env aenv. TypeR a -> OpenFun env aenv (a -> a)
identity (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
ShapeR sh
shR)) (TypeR e -> Fun aenv (e -> e)
forall a env aenv. TypeR a -> OpenFun env aenv (a -> a)
identity TypeR e
TypeR e
tR) Var ArrayR aenv (Array sh e)
v

-- Get the shape of a delayed array
--
shape :: HasCallStack => Cunctation r aenv (Array sh e) -> Exp aenv sh
shape :: forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Exp aenv sh
shape Cunctation r aenv (Array sh e)
cc =
  case Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
delaying Cunctation r aenv (Array sh e)
cc of
    Step ArrayR (Array sh' b)
_ Exp aenv sh'
sh Fun aenv (sh' -> sh)
_ Fun aenv (a -> b)
_ ArrayVar aenv (Array sh a)
_ -> Exp aenv sh
Exp aenv sh'
sh
    Yield ArrayR (Array sh e)
_ Exp aenv sh
sh Fun aenv (sh -> e)
_    -> Exp aenv sh
Exp aenv sh
sh


-- prjExtend :: Kit acc => Extend acc env env' -> Idx env' t -> PreOpenAcc acc env' t
-- prjExtend (PushEnv _   v) ZeroIdx       = weakenA rebuildAcc SuccIdx v
-- prjExtend (PushEnv env _) (SuccIdx idx) = weakenA rebuildAcc SuccIdx $ prjExtend env idx
-- prjExtend _               _             = $internalError "prjExtend" "inconsistent valuation"

{--
-- Rearrange type arguments to fit with Sink type class.
newtype SinkSeq acc senv aenv a = SinkSeq { unSinkSeq :: PreOpenSeq acc aenv senv a }

-- sink for sequences.
sinkSeq :: Kit acc => Extend acc aenv aenv' -> PreOpenSeq acc aenv senv a -> PreOpenSeq acc aenv' senv a
sinkSeq env s = unSinkSeq $ sink env (SinkSeq s)

instance Kit acc => Sink (SinkSeq acc senv) where
  weaken :: forall aenv aenv' arrs. aenv :> aenv' -> SinkSeq acc senv aenv arrs -> SinkSeq acc senv aenv' arrs
  weaken k (SinkSeq s) = SinkSeq $
    case s of
      Producer p s' -> Producer   (weakenP p) (weakenL s')
      Consumer c    -> Consumer   (weakenC c)
      Reify ix      -> Reify      ix

    where
      weakenL :: forall senv' arrs'. PreOpenSeq acc aenv senv' arrs' -> PreOpenSeq acc aenv' senv' arrs'
      weakenL s' = unSinkSeq (weaken k (SinkSeq s'))

      weakenP :: forall a. Producer acc aenv senv a -> Producer acc aenv' senv a
      weakenP p =
        case p of
          StreamIn arrs        -> StreamIn arrs
          ToSeq slix sh a      -> ToSeq slix sh (weaken k a)
          MapSeq f x           -> MapSeq (weaken k f) x
          ChunkedMapSeq f x    -> ChunkedMapSeq (weaken k f) x
          ZipWithSeq f x y     -> ZipWithSeq (weaken k f) x y
          ScanSeq f a x        -> ScanSeq (weaken k f) (weaken k a) x

      weakenC :: forall a. Consumer acc aenv senv a -> Consumer acc aenv' senv a
      weakenC c =
        case c of
          FoldSeq f a x        -> FoldSeq (weaken k f) (weaken k a) x
          FoldSeqFlatten f a x -> FoldSeqFlatten (weaken k f) (weaken k a) x
          Stuple t             ->
            let wk :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv' senv) t
                wk NilAtup        = NilAtup
                wk (SnocAtup t c) = wk t `SnocAtup` weakenC c
            in
            Stuple (wk t)
--}

-- Array fusion of a de Bruijn computation AST
-- ===========================================

-- Array computations
-- ------------------

-- Evaluate a delayed computation and tie the recursive knot
--
-- We do a bit of extra work to (try to) maintain that terms should be left
-- at their lowest common use site. SEE: [Fusion and the lowest common use site]
--
computeAcc
    :: HasCallStack
    => Embed OpenAcc aenv arrs
    -> OpenAcc       aenv arrs
computeAcc :: forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Embed      Extend ArrayR OpenAcc aenv aenv'
BaseEnv              Cunctation r aenv' arrs
cc) = PreOpenAcc OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute Cunctation r aenv arrs
Cunctation r aenv' arrs
cc)
computeAcc (Embed env :: Extend ArrayR OpenAcc aenv aenv'
env@(PushEnv Extend ArrayR OpenAcc aenv env'1
bot LeftHandSide ArrayR t env'1 aenv'
lhs OpenAcc env'1 t
top) Cunctation r aenv' arrs
cc) =
  Cunctation r aenv' arrs
-> Either (Cunctation M aenv' arrs) (Cunctation D aenv' arrs)
forall r aenv a.
HasCallStack =>
Cunctation r aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
simplify Cunctation r aenv' arrs
cc Either (Cunctation M aenv' arrs) (Cunctation D aenv' arrs)
-> (Either (Cunctation M aenv' arrs) (Cunctation D aenv' arrs)
    -> OpenAcc aenv arrs)
-> OpenAcc aenv arrs
forall a b. a -> (a -> b) -> b
& \case
    Left (Done ArrayVars aenv' arrs
v) -> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv aenv'
env ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVars aenv' arrs -> OpenAcc aenv' arrs
forall (acc :: * -> * -> *) aenv arrs.
InjectAcc acc -> ArrayVars aenv arrs -> acc aenv arrs
avarsIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVars aenv' arrs
v)
    Right Cunctation D aenv' arrs
d       -> Cunctation D aenv' arrs
d Cunctation D aenv' arrs
-> (Cunctation D aenv' arrs -> OpenAcc aenv arrs)
-> OpenAcc aenv arrs
forall a b. a -> (a -> b) -> b
& \case
      Yield ArrayR (Array sh e)
repr Exp aenv' sh
sh Fun aenv' (sh -> e)
f
        -> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv aenv'
env (PreOpenAcc OpenAcc aenv' arrs -> OpenAcc aenv' arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (ArrayR (Array sh e)
-> Exp aenv' sh
-> Fun aenv' (sh -> e)
-> PreOpenAcc OpenAcc aenv' (Array sh e)
forall sh e aenv (acc :: * -> * -> *).
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> PreOpenAcc acc aenv (Array sh e)
Generate ArrayR (Array sh e)
repr Exp aenv' sh
sh Fun aenv' (sh -> e)
f))

      Step ArrayR (Array sh' b)
repr Exp aenv' sh'
sh Fun aenv' (sh' -> sh)
p Fun aenv' (a -> b)
f v :: ArrayVar aenv' (Array sh a)
v@(Var ArrayR (Array sh a)
_ Idx aenv' (Array sh a)
ix)
        | Just sh' :~: sh
Refl <- Exp aenv' sh' -> OpenExp () aenv' sh -> Maybe (sh' :~: sh)
forall env aenv s t.
OpenExp env aenv s -> OpenExp env aenv t -> Maybe (s :~: t)
matchOpenExp Exp aenv' sh'
sh (ArrayVar aenv' (Array sh a) -> OpenExp () aenv' sh
forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape ArrayVar aenv' (Array sh a)
v)
        , Just sh' :~: sh
Refl <- Fun aenv' (sh' -> sh) -> Maybe (sh' :~: sh)
forall env aenv a b. OpenFun env aenv (a -> b) -> Maybe (a :~: b)
isIdentity Fun aenv' (sh' -> sh)
p
        -> case Idx aenv' (Array sh a)
ix of
             Idx aenv' (Array sh a)
ZeroIdx
               | LeftHandSideSingle ArrayR{} <- LeftHandSide ArrayR t env'1 aenv'
lhs
               , Just (OpenAccFun OpenFun () env'1 (a -> b)
g) <- (aenv' :?> env'1)
-> OpenAccFun Any () aenv' (a -> b)
-> Maybe (OpenAccFun Any () env'1 (a -> b))
forall (f :: * -> * -> *) env env' t.
Rebuildable f =>
(env :?> env') -> f env t -> Maybe (f env' t)
strengthen Idx aenv' t' -> Maybe (Idx env'1 t')
Idx (env'1, Array sh' a) t' -> Maybe (Idx env'1 t')
aenv' :?> env'1
forall aenv a t'. Idx (aenv, a) t' -> Maybe (Idx aenv t')
noTop (Fun aenv' (a -> b) -> OpenAccFun Any () aenv' (a -> b)
forall env aenv a (acc :: * -> * -> *).
OpenFun env aenv a -> OpenAccFun acc env aenv a
OpenAccFun Fun aenv' (a -> b)
f)
                    -> Extend ArrayR OpenAcc aenv env'1
-> OpenAcc env'1 arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv env'1
bot (PreOpenAcc OpenAcc env'1 arrs -> OpenAcc env'1 arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (TypeR b
-> OpenFun () env'1 (a -> b)
-> OpenAcc env'1 (Array sh' a)
-> PreOpenAcc OpenAcc env'1 (Array sh' b)
forall e' aenv e (acc :: * -> * -> *) sh.
TypeR e'
-> Fun aenv (e -> e')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Map (ArrayR (Array sh' b) -> TypeR b
forall sh e. ArrayR (Array sh e) -> TypeR e
arrayRtype ArrayR (Array sh' b)
repr) OpenFun () env'1 (a -> b)
g OpenAcc env'1 t
OpenAcc env'1 (Array sh' a)
top))
             Idx aenv' (Array sh a)
_      -> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv aenv'
env (PreOpenAcc OpenAcc aenv' arrs -> OpenAcc aenv' arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (TypeR b
-> Fun aenv' (a -> b)
-> OpenAcc aenv' (Array sh a)
-> PreOpenAcc OpenAcc aenv' (Array sh b)
forall e' aenv e (acc :: * -> * -> *) sh.
TypeR e'
-> Fun aenv (e -> e')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Map (ArrayR (Array sh' b) -> TypeR b
forall sh e. ArrayR (Array sh e) -> TypeR e
arrayRtype ArrayR (Array sh' b)
repr) Fun aenv' (a -> b)
f ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVar aenv' (Array sh a) -> OpenAcc aenv' (Array sh a)
forall (acc :: * -> * -> *) aenv a.
InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVar aenv' (Array sh a)
v)))

        | Just a :~: b
Refl <- Fun aenv' (a -> b) -> Maybe (a :~: b)
forall env aenv a b. OpenFun env aenv (a -> b) -> Maybe (a :~: b)
isIdentity Fun aenv' (a -> b)
f
        -> case Idx aenv' (Array sh a)
ix of
             Idx aenv' (Array sh a)
ZeroIdx
               | LeftHandSideSingle ArrayR{} <- LeftHandSide ArrayR t env'1 aenv'
lhs
               , Just (OpenAccFun OpenFun () env'1 (sh' -> sh)
q)  <- (aenv' :?> env'1)
-> OpenAccFun Any () aenv' (sh' -> sh)
-> Maybe (OpenAccFun Any () env'1 (sh' -> sh))
forall (f :: * -> * -> *) env env' t.
Rebuildable f =>
(env :?> env') -> f env t -> Maybe (f env' t)
strengthen Idx aenv' t' -> Maybe (Idx env'1 t')
Idx (env'1, Array sh a) t' -> Maybe (Idx env'1 t')
aenv' :?> env'1
forall aenv a t'. Idx (aenv, a) t' -> Maybe (Idx aenv t')
noTop (Fun aenv' (sh' -> sh) -> OpenAccFun Any () aenv' (sh' -> sh)
forall env aenv a (acc :: * -> * -> *).
OpenFun env aenv a -> OpenAccFun acc env aenv a
OpenAccFun Fun aenv' (sh' -> sh)
p)
               , Just (OpenAccExp OpenExp () env'1 sh'
sz) <- (aenv' :?> env'1)
-> OpenAccExp Any () aenv' sh'
-> Maybe (OpenAccExp Any () env'1 sh')
forall (f :: * -> * -> *) env env' t.
Rebuildable f =>
(env :?> env') -> f env t -> Maybe (f env' t)
strengthen Idx aenv' t' -> Maybe (Idx env'1 t')
Idx (env'1, Array sh a) t' -> Maybe (Idx env'1 t')
aenv' :?> env'1
forall aenv a t'. Idx (aenv, a) t' -> Maybe (Idx aenv t')
noTop (Exp aenv' sh' -> OpenAccExp Any () aenv' sh'
forall env aenv a (acc :: * -> * -> *).
OpenExp env aenv a -> OpenAccExp acc env aenv a
OpenAccExp Exp aenv' sh'
sh)
                    -> Extend ArrayR OpenAcc aenv env'1
-> OpenAcc env'1 arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv env'1
bot (PreOpenAcc OpenAcc env'1 arrs -> OpenAcc env'1 arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (ShapeR sh'
-> OpenExp () env'1 sh'
-> OpenFun () env'1 (sh' -> sh)
-> OpenAcc env'1 (Array sh a)
-> PreOpenAcc OpenAcc env'1 (Array sh' a)
forall sh' aenv sh (acc :: * -> * -> *) e.
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Backpermute (ArrayR (Array sh' b) -> ShapeR sh'
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh' b)
repr) OpenExp () env'1 sh'
sz OpenFun () env'1 (sh' -> sh)
q OpenAcc env'1 t
OpenAcc env'1 (Array sh a)
top))
             Idx aenv' (Array sh a)
_      -> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv aenv'
env (PreOpenAcc OpenAcc aenv' arrs -> OpenAcc aenv' arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (ShapeR sh'
-> Exp aenv' sh'
-> Fun aenv' (sh' -> sh)
-> OpenAcc aenv' (Array sh a)
-> PreOpenAcc OpenAcc aenv' (Array sh' a)
forall sh' aenv sh (acc :: * -> * -> *) e.
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Backpermute (ArrayR (Array sh' b) -> ShapeR sh'
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh' b)
repr) Exp aenv' sh'
sh Fun aenv' (sh' -> sh)
p ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVar aenv' (Array sh a) -> OpenAcc aenv' (Array sh a)
forall (acc :: * -> * -> *) aenv a.
InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVar aenv' (Array sh a)
v)))

        | Bool
otherwise
        -> case Idx aenv' (Array sh a)
ix of
             Idx aenv' (Array sh a)
ZeroIdx
               | LeftHandSideSingle ArrayR{} <- LeftHandSide ArrayR t env'1 aenv'
lhs
               , Just (OpenAccFun OpenFun () env'1 (a -> b)
g)  <- (aenv' :?> env'1)
-> OpenAccFun Any () aenv' (a -> b)
-> Maybe (OpenAccFun Any () env'1 (a -> b))
forall (f :: * -> * -> *) env env' t.
Rebuildable f =>
(env :?> env') -> f env t -> Maybe (f env' t)
strengthen Idx aenv' t' -> Maybe (Idx env'1 t')
Idx (env'1, Array sh a) t' -> Maybe (Idx env'1 t')
aenv' :?> env'1
forall aenv a t'. Idx (aenv, a) t' -> Maybe (Idx aenv t')
noTop (Fun aenv' (a -> b) -> OpenAccFun Any () aenv' (a -> b)
forall env aenv a (acc :: * -> * -> *).
OpenFun env aenv a -> OpenAccFun acc env aenv a
OpenAccFun Fun aenv' (a -> b)
f)
               , Just (OpenAccFun OpenFun () env'1 (sh' -> sh)
q)  <- (aenv' :?> env'1)
-> OpenAccFun Any () aenv' (sh' -> sh)
-> Maybe (OpenAccFun Any () env'1 (sh' -> sh))
forall (f :: * -> * -> *) env env' t.
Rebuildable f =>
(env :?> env') -> f env t -> Maybe (f env' t)
strengthen Idx aenv' t' -> Maybe (Idx env'1 t')
Idx (env'1, Array sh a) t' -> Maybe (Idx env'1 t')
aenv' :?> env'1
forall aenv a t'. Idx (aenv, a) t' -> Maybe (Idx aenv t')
noTop (Fun aenv' (sh' -> sh) -> OpenAccFun Any () aenv' (sh' -> sh)
forall env aenv a (acc :: * -> * -> *).
OpenFun env aenv a -> OpenAccFun acc env aenv a
OpenAccFun Fun aenv' (sh' -> sh)
p)
               , Just (OpenAccExp OpenExp () env'1 sh'
sz) <- (aenv' :?> env'1)
-> OpenAccExp Any () aenv' sh'
-> Maybe (OpenAccExp Any () env'1 sh')
forall (f :: * -> * -> *) env env' t.
Rebuildable f =>
(env :?> env') -> f env t -> Maybe (f env' t)
strengthen Idx aenv' t' -> Maybe (Idx env'1 t')
Idx (env'1, Array sh a) t' -> Maybe (Idx env'1 t')
aenv' :?> env'1
forall aenv a t'. Idx (aenv, a) t' -> Maybe (Idx aenv t')
noTop (Exp aenv' sh' -> OpenAccExp Any () aenv' sh'
forall env aenv a (acc :: * -> * -> *).
OpenExp env aenv a -> OpenAccExp acc env aenv a
OpenAccExp Exp aenv' sh'
sh)
                    -> Extend ArrayR OpenAcc aenv env'1
-> OpenAcc env'1 arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv env'1
bot (PreOpenAcc OpenAcc env'1 arrs -> OpenAcc env'1 arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (ArrayR (Array sh' b)
-> OpenExp () env'1 sh'
-> OpenFun () env'1 (sh' -> sh)
-> OpenFun () env'1 (a -> b)
-> OpenAcc env'1 (Array sh a)
-> PreOpenAcc OpenAcc env'1 (Array sh' b)
forall sh' b aenv sh a1 (acc :: * -> * -> *).
ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> acc aenv (Array sh a1)
-> PreOpenAcc acc aenv (Array sh' b)
Transform ArrayR (Array sh' b)
repr OpenExp () env'1 sh'
sz OpenFun () env'1 (sh' -> sh)
q OpenFun () env'1 (a -> b)
g OpenAcc env'1 t
OpenAcc env'1 (Array sh a)
top))
             Idx aenv' (Array sh a)
_      -> Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' arrs -> OpenAcc aenv arrs
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv aenv'
env (PreOpenAcc OpenAcc aenv' arrs -> OpenAcc aenv' arrs
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (ArrayR (Array sh' b)
-> Exp aenv' sh'
-> Fun aenv' (sh' -> sh)
-> Fun aenv' (a -> b)
-> OpenAcc aenv' (Array sh a)
-> PreOpenAcc OpenAcc aenv' (Array sh' b)
forall sh' b aenv sh a1 (acc :: * -> * -> *).
ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> acc aenv (Array sh a1)
-> PreOpenAcc acc aenv (Array sh' b)
Transform ArrayR (Array sh' b)
repr Exp aenv' sh'
sh Fun aenv' (sh' -> sh)
p Fun aenv' (a -> b)
f ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVar aenv' (Array sh a) -> OpenAcc aenv' (Array sh a)
forall (acc :: * -> * -> *) aenv a.
InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVar aenv' (Array sh a)
v)))

  where
    bindA :: HasCallStack
          => Extend ArrayR OpenAcc aenv aenv'
          -> OpenAcc aenv' a
          -> OpenAcc aenv  a
    bindA :: forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv aenv'
BaseEnv             OpenAcc aenv' a
b = OpenAcc aenv a
OpenAcc aenv' a
b
    bindA (PushEnv Extend ArrayR OpenAcc aenv env'1
env LeftHandSide ArrayR t env'1 aenv'
lhs OpenAcc env'1 t
a) OpenAcc aenv' a
b
      -- If the freshly bound value is directly, returned, we don't have to bind it in a
      -- let. We can do this if the left hand side does not contain wildcards (other than
      -- wildcards for unit / nil) and if the value contains the same variables.
      | Just ArrayVars aenv' a
vars <- OpenAcc aenv' a -> Maybe (ArrayVars aenv' a)
forall aenv a. OpenAcc aenv a -> Maybe (ArrayVars aenv a)
extractOpenArrayVars OpenAcc aenv' a
b
      , Just t :~: a
Refl <- LeftHandSide ArrayR t env'1 aenv'
-> ArrayVars aenv' a -> Maybe (t :~: a)
forall (s :: * -> *) a env1 env2 b.
LeftHandSide s a env1 env2 -> Vars s env2 b -> Maybe (a :~: b)
bindingIsTrivial LeftHandSide ArrayR t env'1 aenv'
lhs ArrayVars aenv' a
vars = Extend ArrayR OpenAcc aenv env'1
-> OpenAcc env'1 a -> OpenAcc aenv a
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv env'1
env OpenAcc env'1 a
OpenAcc env'1 t
a
      | Bool
otherwise                              = Extend ArrayR OpenAcc aenv env'1
-> OpenAcc env'1 a -> OpenAcc aenv a
forall aenv aenv' a.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' a -> OpenAcc aenv a
bindA Extend ArrayR OpenAcc aenv env'1
env (PreOpenAcc OpenAcc env'1 a -> OpenAcc env'1 a
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (LeftHandSide ArrayR t env'1 aenv'
-> OpenAcc env'1 t -> OpenAcc aenv' a -> PreOpenAcc OpenAcc env'1 a
forall bndArrs aenv aenv' (acc :: * -> * -> *) a.
ALeftHandSide bndArrs aenv aenv'
-> acc aenv bndArrs -> acc aenv' a -> PreOpenAcc acc aenv a
Alet LeftHandSide ArrayR t env'1 aenv'
lhs OpenAcc env'1 t
a OpenAcc aenv' a
b))

    noTop :: (aenv, a) :?> aenv
    noTop :: forall aenv a t'. Idx (aenv, a) t' -> Maybe (Idx aenv t')
noTop Idx (aenv, a) t'
ZeroIdx      = Maybe (Idx aenv t')
forall a. Maybe a
Nothing
    noTop (SuccIdx Idx env t'
ix) = Idx aenv t' -> Maybe (Idx aenv t')
forall a. a -> Maybe a
Just Idx aenv t'
Idx env t'
ix



-- Convert the internal representation of delayed arrays into a real AST
-- node. Use the most specific version of a combinator whenever possible.
--
compute
    :: HasCallStack
    => Cunctation r       aenv arrs
    -> PreOpenAcc OpenAcc aenv arrs
compute :: forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute Cunctation r aenv arrs
cc = Cunctation r aenv arrs
-> Either (Cunctation M aenv arrs) (Cunctation D aenv arrs)
forall r aenv a.
HasCallStack =>
Cunctation r aenv a
-> Either (Cunctation M aenv a) (Cunctation D aenv a)
simplify Cunctation r aenv arrs
cc Either (Cunctation M aenv arrs) (Cunctation D aenv arrs)
-> (Either (Cunctation M aenv arrs) (Cunctation D aenv arrs)
    -> PreOpenAcc OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs
forall a b. a -> (a -> b) -> b
& \case
  Left (Done ArrayVars aenv arrs
v) -> ArrayVars aenv arrs
v ArrayVars aenv arrs
-> (ArrayVars aenv arrs -> PreOpenAcc OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs
forall a b. a -> (a -> b) -> b
& \case
    ArrayVars aenv arrs
TupRunit                      -> PreOpenAcc OpenAcc aenv arrs
PreOpenAcc OpenAcc aenv ()
forall (acc :: * -> * -> *) aenv. PreOpenAcc acc aenv ()
Anil
    TupRsingle v :: Var ArrayR aenv arrs
v@(Var ArrayR{} Idx aenv arrs
_) -> ArrayVar aenv (Array sh e) -> PreOpenAcc OpenAcc aenv (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar Var ArrayR aenv arrs
ArrayVar aenv (Array sh e)
v
    TupRpair TupR (Var ArrayR aenv) a1
v1 TupR (Var ArrayR aenv) b
v2                -> (forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> TupR (Var ArrayR aenv) a1 -> OpenAcc aenv a1
forall (acc :: * -> * -> *) aenv arrs.
InjectAcc acc -> ArrayVars aenv arrs -> acc aenv arrs
avarsIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc TupR (Var ArrayR aenv) a1
v1 OpenAcc aenv a1
-> OpenAcc aenv b -> PreOpenAcc OpenAcc aenv (a1, b)
forall (acc :: * -> * -> *) aenv as bs.
acc aenv as -> acc aenv bs -> PreOpenAcc acc aenv (as, bs)
`Apair` (forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> TupR (Var ArrayR aenv) b -> OpenAcc aenv b
forall (acc :: * -> * -> *) aenv arrs.
InjectAcc acc -> ArrayVars aenv arrs -> acc aenv arrs
avarsIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc TupR (Var ArrayR aenv) b
v2
  Right Cunctation D aenv arrs
d -> Cunctation D aenv arrs
d Cunctation D aenv arrs
-> (Cunctation D aenv arrs -> PreOpenAcc OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs
forall a b. a -> (a -> b) -> b
& \case
    Yield ArrayR (Array sh e)
aR Exp aenv sh
sh Fun aenv (sh -> e)
f                 -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> PreOpenAcc OpenAcc aenv (Array sh e)
forall sh e aenv (acc :: * -> * -> *).
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> PreOpenAcc acc aenv (Array sh e)
Generate ArrayR (Array sh e)
aR Exp aenv sh
sh Fun aenv (sh -> e)
f
    Step (ArrayR ShapeR sh
shR TypeR e
tR) Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a -> b)
f ArrayVar aenv (Array sh a)
v
      | Just sh' :~: sh
Refl <- Exp aenv sh' -> OpenExp () aenv sh -> Maybe (sh' :~: sh)
forall env aenv s t.
OpenExp env aenv s -> OpenExp env aenv t -> Maybe (s :~: t)
matchOpenExp Exp aenv sh'
sh (ArrayVar aenv (Array sh a) -> OpenExp () aenv sh
forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape ArrayVar aenv (Array sh a)
v)
      , Just sh' :~: sh
Refl <- Fun aenv (sh' -> sh) -> Maybe (sh' :~: sh)
forall env aenv a b. OpenFun env aenv (a -> b) -> Maybe (a :~: b)
isIdentity Fun aenv (sh' -> sh)
p -> TypeR e
-> Fun aenv (a -> e)
-> OpenAcc aenv (Array sh a)
-> PreOpenAcc OpenAcc aenv (Array sh e)
forall e' aenv e (acc :: * -> * -> *) sh.
TypeR e'
-> Fun aenv (e -> e')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Map TypeR e
tR Fun aenv (a -> b)
Fun aenv (a -> e)
f ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVar aenv (Array sh a) -> OpenAcc aenv (Array sh a)
forall (acc :: * -> * -> *) aenv a.
InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVar aenv (Array sh a)
v)
      | Just a :~: b
Refl <- Fun aenv (a -> b) -> Maybe (a :~: b)
forall env aenv a b. OpenFun env aenv (a -> b) -> Maybe (a :~: b)
isIdentity Fun aenv (a -> b)
f -> ShapeR sh
-> Exp aenv sh
-> Fun aenv (sh -> sh)
-> OpenAcc aenv (Array sh a)
-> PreOpenAcc OpenAcc aenv (Array sh a)
forall sh' aenv sh (acc :: * -> * -> *) e.
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Backpermute ShapeR sh
shR Exp aenv sh'
Exp aenv sh
sh Fun aenv (sh' -> sh)
Fun aenv (sh -> sh)
p ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVar aenv (Array sh a) -> OpenAcc aenv (Array sh a)
forall (acc :: * -> * -> *) aenv a.
InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVar aenv (Array sh a)
v)
      | Bool
otherwise                 -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> sh)
-> Fun aenv (a -> e)
-> OpenAcc aenv (Array sh a)
-> PreOpenAcc OpenAcc aenv (Array sh e)
forall sh' b aenv sh a1 (acc :: * -> * -> *).
ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> acc aenv (Array sh a1)
-> PreOpenAcc acc aenv (Array sh' b)
Transform (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shR TypeR e
tR) Exp aenv sh'
Exp aenv sh
sh Fun aenv (sh' -> sh)
Fun aenv (sh -> sh)
p Fun aenv (a -> b)
Fun aenv (a -> e)
f ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVar aenv (Array sh a) -> OpenAcc aenv (Array sh a)
forall (acc :: * -> * -> *) aenv a.
InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVar aenv (Array sh a)
v)


-- Representation of a generator as a delayed array
--
generateD
    :: HasCallStack
    => ArrayR (Array sh e)
    -> Exp aenv sh
    -> Fun aenv (sh -> e)
    -> Embed OpenAcc aenv (Array sh e)
generateD :: forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Embed OpenAcc aenv (Array sh e)
generateD ArrayR (Array sh e)
repr Exp aenv sh
sh Fun aenv (sh -> e)
f
  = Text
-> Embed OpenAcc aenv (Array sh e)
-> Embed OpenAcc aenv (Array sh e)
forall a. Text -> a -> a
Stats.ruleFired Text
"generateD"
  (Embed OpenAcc aenv (Array sh e)
 -> Embed OpenAcc aenv (Array sh e))
-> Embed OpenAcc aenv (Array sh e)
-> Embed OpenAcc aenv (Array sh e)
forall a b. (a -> b) -> a -> b
$ Extend ArrayR OpenAcc aenv aenv
-> Cunctation D aenv (Array sh e)
-> Embed OpenAcc aenv (Array sh e)
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv
forall (s :: * -> *) (f :: * -> * -> *) env. Extend s f env env
BaseEnv (ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> Cunctation D aenv (Array sh e)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield ArrayR (Array sh e)
repr Exp aenv sh
sh Fun aenv (sh -> e)
f)


-- Fuse a unary function into a delayed array. Also looks for unzips which can
-- be executed in constant time; SEE [unzipD]
--
mapD :: HasCallStack
     => TypeR b
     -> Fun           aenv (a -> b)
     -> Embed OpenAcc aenv (Array sh a)
     -> Embed OpenAcc aenv (Array sh b)
mapD :: forall b aenv a sh.
HasCallStack =>
TypeR b
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh b)
mapD TypeR b
tR Fun aenv (a -> b)
f (TypeR b
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Maybe (Embed OpenAcc aenv (Array sh b))
forall b aenv a sh.
HasCallStack =>
TypeR b
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Maybe (Embed OpenAcc aenv (Array sh b))
unzipD TypeR b
tR Fun aenv (a -> b)
f -> Just Embed OpenAcc aenv (Array sh b)
a) = Embed OpenAcc aenv (Array sh b)
a
mapD TypeR b
tR Fun aenv (a -> b)
f (Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' (Array sh a)
cc)
  = Text
-> Embed OpenAcc aenv (Array sh b)
-> Embed OpenAcc aenv (Array sh b)
forall a. Text -> a -> a
Stats.ruleFired Text
"mapD"
  (Embed OpenAcc aenv (Array sh b)
 -> Embed OpenAcc aenv (Array sh b))
-> Embed OpenAcc aenv (Array sh b)
-> Embed OpenAcc aenv (Array sh b)
forall a b. (a -> b) -> a -> b
$ Extend ArrayR OpenAcc aenv aenv'
-> Cunctation D aenv' (Array sh b)
-> Embed OpenAcc aenv (Array sh b)
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env (Cunctation D aenv' (Array sh a) -> Cunctation D aenv' (Array sh b)
go (Cunctation r aenv' (Array sh a) -> Cunctation D aenv' (Array sh a)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
delaying Cunctation r aenv' (Array sh a)
cc))
  where
    go :: Cunctation D aenv' (Array sh a) -> Cunctation D aenv' (Array sh b)
go (Step (ArrayR ShapeR sh
shR TypeR e
_) Exp aenv' sh'
sh Fun aenv' (sh' -> sh)
ix Fun aenv' (a -> b)
g ArrayVar aenv' (Array sh a)
v) = ArrayR (Array sh b)
-> Exp aenv' sh
-> Fun aenv' (sh -> sh)
-> Fun aenv' (a -> b)
-> ArrayVar aenv' (Array sh a)
-> Cunctation D aenv' (Array sh b)
forall env' e aenv sh a.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> sh)
-> Fun aenv (a -> e)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array env' e)
Step  (ShapeR sh -> TypeR b -> ArrayR (Array sh b)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shR TypeR b
tR) Exp aenv' sh'
Exp aenv' sh
sh Fun aenv' (sh' -> sh)
Fun aenv' (sh -> sh)
ix (Extend ArrayR OpenAcc aenv aenv'
-> Fun aenv (a -> b) -> OpenFun () aenv' (a -> b)
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv aenv'
env Fun aenv (a -> b)
f OpenFun () aenv' (a -> b)
-> OpenFun () aenv' (a -> a) -> Fun aenv' (a -> b)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` OpenFun () aenv' (a -> a)
Fun aenv' (a -> b)
g) ArrayVar aenv' (Array sh a)
v
    go (Yield (ArrayR ShapeR sh
shR TypeR e
_) Exp aenv' sh
sh Fun aenv' (sh -> e)
g)     = ArrayR (Array sh b)
-> Exp aenv' sh
-> Fun aenv' (sh -> b)
-> Cunctation D aenv' (Array sh b)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield (ShapeR sh -> TypeR b -> ArrayR (Array sh b)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shR TypeR b
tR) Exp aenv' sh
Exp aenv' sh
sh    (Extend ArrayR OpenAcc aenv aenv'
-> Fun aenv (a -> b) -> OpenFun () aenv' (a -> b)
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv aenv'
env Fun aenv (a -> b)
f OpenFun () aenv' (a -> b)
-> OpenFun () aenv' (sh -> a) -> Fun aenv' (sh -> b)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` Fun aenv' (sh -> e)
OpenFun () aenv' (sh -> a)
g)


-- If we are unzipping a manifest array then force the term to be computed;
-- a backend will be able to execute this in constant time.
--
unzipD
    :: HasCallStack
    => TypeR b
    -> Fun                  aenv (a -> b)
    -> Embed OpenAcc        aenv (Array sh a)
    -> Maybe (Embed OpenAcc aenv (Array sh b))
unzipD :: forall b aenv a sh.
HasCallStack =>
TypeR b
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Maybe (Embed OpenAcc aenv (Array sh b))
unzipD TypeR b
tR Fun aenv (a -> b)
f (Embed Extend ArrayR OpenAcc aenv aenv'
env cc :: Cunctation r aenv' (Array sh a)
cc@(Done ArrayVars aenv' (Array sh a)
v))
  | Lam ELeftHandSide a () env'
lhs (Body OpenExp env' aenv t1
a) <- Fun aenv (a -> b)
f
  , Just ExpVars env' t1
vars        <- OpenExp env' aenv t1 -> Maybe (ExpVars env' t1)
forall env aenv a. OpenExp env aenv a -> Maybe (ExpVars env a)
extractExpVars OpenExp env' aenv t1
a
  , ArrayR ShapeR sh
shR TypeR e
_     <- Cunctation r aenv' (Array sh a) -> ArrayR (Array sh a)
forall (f :: * -> * -> *) aenv sh e.
HasArraysR f =>
f aenv (Array sh e) -> ArrayR (Array sh e)
arrayR Cunctation r aenv' (Array sh a)
cc
  , OpenFun () aenv' (a -> t1)
f'               <- ELeftHandSide a () env'
-> OpenFun env' aenv' t1 -> OpenFun () aenv' (a -> t1)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam ELeftHandSide a () env'
lhs (OpenFun env' aenv' t1 -> OpenFun () aenv' (a -> t1))
-> OpenFun env' aenv' t1 -> OpenFun () aenv' (a -> t1)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv' t1 -> OpenFun env' aenv' t1
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv' t1 -> OpenFun env' aenv' t1)
-> OpenExp env' aenv' t1 -> OpenFun env' aenv' t1
forall a b. (a -> b) -> a -> b
$ ExpVars env' t1 -> OpenExp env' aenv' t1
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars ExpVars env' t1
vars
  = Embed OpenAcc aenv (Array sh b)
-> Maybe (Embed OpenAcc aenv (Array sh b))
forall a. a -> Maybe a
Just (Embed OpenAcc aenv (Array sh b)
 -> Maybe (Embed OpenAcc aenv (Array sh b)))
-> Embed OpenAcc aenv (Array sh b)
-> Maybe (Embed OpenAcc aenv (Array sh b))
forall a b. (a -> b) -> a -> b
$ Extend ArrayR OpenAcc aenv (aenv', Array sh b)
-> Cunctation M (aenv', Array sh b) (Array sh b)
-> Embed OpenAcc aenv (Array sh b)
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
env Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array sh b)
-> Extend ArrayR OpenAcc aenv (aenv', Array sh b)
forall (acc :: * -> * -> *) aenv aenv' sh e.
HasArraysR acc =>
Extend ArrayR acc aenv aenv'
-> acc aenv' (Array sh e)
-> Extend ArrayR acc aenv (aenv', Array sh e)
`pushArrayEnv` PreOpenAcc OpenAcc aenv' (Array sh b) -> OpenAcc aenv' (Array sh b)
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (TypeR b
-> Fun aenv' (a -> b)
-> OpenAcc aenv' (Array sh a)
-> PreOpenAcc OpenAcc aenv' (Array sh b)
forall e' aenv e (acc :: * -> * -> *) sh.
TypeR e'
-> Fun aenv (e -> e')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Map TypeR b
tR Fun aenv' (a -> b)
OpenFun () aenv' (a -> t1)
f' (OpenAcc aenv' (Array sh a)
 -> PreOpenAcc OpenAcc aenv' (Array sh b))
-> OpenAcc aenv' (Array sh a)
-> PreOpenAcc OpenAcc aenv' (Array sh b)
forall a b. (a -> b) -> a -> b
$ (forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVars aenv' (Array sh a) -> OpenAcc aenv' (Array sh a)
forall (acc :: * -> * -> *) aenv arrs.
InjectAcc acc -> ArrayVars aenv arrs -> acc aenv arrs
avarsIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVars aenv' (Array sh a)
ArrayVars aenv' (Array sh a)
v)) (Cunctation M (aenv', Array sh b) (Array sh b)
 -> Embed OpenAcc aenv (Array sh b))
-> Cunctation M (aenv', Array sh b) (Array sh b)
-> Embed OpenAcc aenv (Array sh b)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh b)
-> Cunctation M (aenv', Array sh b) (Array sh b)
forall sh e aenv.
ArrayR (Array sh e) -> Cunctation M (aenv, Array sh e) (Array sh e)
doneZeroIdx (ArrayR (Array sh b)
 -> Cunctation M (aenv', Array sh b) (Array sh b))
-> ArrayR (Array sh b)
-> Cunctation M (aenv', Array sh b) (Array sh b)
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> TypeR b -> ArrayR (Array sh b)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh
shR TypeR b
tR
unzipD TypeR b
_ Fun aenv (a -> b)
_ Embed OpenAcc aenv (Array sh a)
_
  = Maybe (Embed OpenAcc aenv (Array sh b))
forall a. Maybe a
Nothing

-- Fuse an index space transformation function that specifies where elements in
-- the destination array read there data from in the source array.
--
backpermuteD
    :: HasCallStack
    => ShapeR sh'
    -> Exp          aenv sh'
    -> Fun          aenv (sh' -> sh)
    -> Cunctation r aenv (Array sh  e)
    -> Cunctation D aenv (Array sh' e)
backpermuteD :: forall sh' aenv sh r e.
HasCallStack =>
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
backpermuteD ShapeR sh'
shR' Exp aenv sh'
sh' Fun aenv (sh' -> sh)
p = Text
-> Cunctation D aenv (Array sh' e)
-> Cunctation D aenv (Array sh' e)
forall a. Text -> a -> a
Stats.ruleFired Text
"backpermuteD" (Cunctation D aenv (Array sh' e)
 -> Cunctation D aenv (Array sh' e))
-> (Cunctation r aenv (Array sh e)
    -> Cunctation D aenv (Array sh' e))
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cunctation D aenv (Array sh e) -> Cunctation D aenv (Array sh' e)
go (Cunctation D aenv (Array sh e) -> Cunctation D aenv (Array sh' e))
-> (Cunctation r aenv (Array sh e)
    -> Cunctation D aenv (Array sh e))
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
delaying
  where
    go :: Cunctation D aenv (Array sh e) -> Cunctation D aenv (Array sh' e)
go (Step (ArrayR ShapeR sh
_ TypeR e
tR) Exp aenv sh'
_ Fun aenv (sh' -> sh)
q Fun aenv (a -> b)
f ArrayVar aenv (Array sh a)
v) = ArrayR (Array sh' e)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a -> e)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array sh' e)
forall env' e aenv sh a.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> sh)
-> Fun aenv (a -> e)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array env' e)
Step  (ShapeR sh' -> TypeR e -> ArrayR (Array sh' e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh'
shR' TypeR e
tR) Exp aenv sh'
sh' (Fun aenv (sh' -> sh)
q Fun aenv (sh' -> sh)
-> OpenFun () aenv (sh' -> sh') -> Fun aenv (sh' -> sh)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` Fun aenv (sh' -> sh)
OpenFun () aenv (sh' -> sh')
p) Fun aenv (a -> b)
Fun aenv (a -> e)
f ArrayVar aenv (Array sh a)
v
    go (Yield (ArrayR ShapeR sh
_ TypeR e
tR) Exp aenv sh
_ Fun aenv (sh -> e)
g)    = ArrayR (Array sh' e)
-> Exp aenv sh'
-> Fun aenv (sh' -> e)
-> Cunctation D aenv (Array sh' e)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield (ShapeR sh' -> TypeR e -> ArrayR (Array sh' e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh'
shR' TypeR e
tR) Exp aenv sh'
sh' (Fun aenv (sh -> e)
OpenFun () aenv (sh -> e)
g OpenFun () aenv (sh -> e)
-> OpenFun () aenv (sh' -> sh) -> Fun aenv (sh' -> e)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` Fun aenv (sh' -> sh)
OpenFun () aenv (sh' -> sh)
p)


-- Transform as a combined map and backwards permutation
--
transformD
    :: HasCallStack
    => ArrayR (Array sh' b)
    -> Exp           aenv sh'
    -> Fun           aenv (sh' -> sh)
    -> Fun           aenv (a   -> b)
    -> Embed OpenAcc aenv (Array sh  a)
    -> Embed OpenAcc aenv (Array sh' b)
transformD :: forall sh' b aenv sh a.
HasCallStack =>
ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh' b)
transformD (ArrayR ShapeR sh
shR' TypeR e
tR) Exp aenv sh'
sh' Fun aenv (sh' -> sh)
p Fun aenv (a -> b)
f
  = Text
-> Embed OpenAcc aenv (Array sh' b)
-> Embed OpenAcc aenv (Array sh' b)
forall a. Text -> a -> a
Stats.ruleFired Text
"transformD"
  (Embed OpenAcc aenv (Array sh' b)
 -> Embed OpenAcc aenv (Array sh' b))
-> (Embed OpenAcc aenv (Array sh a)
    -> Embed OpenAcc aenv (Array sh' b))
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh' b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' (Array sh e)
 -> Cunctation D aenv' (Array sh' b))
-> Embed OpenAcc aenv (Array sh e)
-> Embed OpenAcc aenv (Array sh' b)
forall aenv as bs.
HasCallStack =>
(forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as -> Cunctation D aenv' bs)
-> Embed OpenAcc aenv as -> Embed OpenAcc aenv bs
fuse ((OpenExp () aenv' sh'
 -> OpenFun () aenv' (sh' -> sh)
 -> Cunctation r aenv' (Array sh e)
 -> Cunctation D aenv' (Array sh' b))
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> Cunctation D aenv' (Array sh' b)
forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env.
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> f2 env' b -> c)
-> f1 env a -> f2 env b -> Extend ArrayR OpenAcc env env' -> c
into2 (ShapeR sh'
-> OpenExp () aenv' sh'
-> OpenFun () aenv' (sh' -> sh)
-> Cunctation r aenv' (Array sh b)
-> Cunctation D aenv' (Array sh' b)
forall sh' aenv sh r e.
HasCallStack =>
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
backpermuteD ShapeR sh'
ShapeR sh
shR') Exp aenv sh'
sh' Fun aenv (sh' -> sh)
p)
  (Embed OpenAcc aenv (Array sh e)
 -> Embed OpenAcc aenv (Array sh' b))
-> (Embed OpenAcc aenv (Array sh a)
    -> Embed OpenAcc aenv (Array sh e))
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh' b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeR e
-> Fun aenv (a -> e)
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh e)
forall b aenv a sh.
HasCallStack =>
TypeR b
-> Fun aenv (a -> b)
-> Embed OpenAcc aenv (Array sh a)
-> Embed OpenAcc aenv (Array sh b)
mapD TypeR e
tR Fun aenv (a -> b)
Fun aenv (a -> e)
f
  where
    fuse :: HasCallStack
         => (forall r aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation D aenv' bs)
         -> Embed OpenAcc aenv as
         -> Embed OpenAcc aenv bs
    fuse :: forall aenv as bs.
HasCallStack =>
(forall r aenv'.
 Extend ArrayR OpenAcc aenv aenv'
 -> Cunctation r aenv' as -> Cunctation D aenv' bs)
-> Embed OpenAcc aenv as -> Embed OpenAcc aenv bs
fuse forall r aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Cunctation D aenv' bs
op (Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' as
cc) = Extend ArrayR OpenAcc aenv aenv'
-> Cunctation D aenv' bs -> Embed OpenAcc aenv bs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Cunctation D aenv' bs
forall r aenv'.
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' as -> Cunctation D aenv' bs
op Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' as
cc)

    into2 :: (HasCallStack, Sink f1, Sink f2)
          => (f1 env' a -> f2 env' b -> c)
          -> f1 env a
          -> f2 env b
          -> Extend ArrayR OpenAcc env env'
          -> c
    into2 :: forall (f1 :: * -> * -> *) (f2 :: * -> * -> *) env' a b c env.
(HasCallStack, Sink f1, Sink f2) =>
(f1 env' a -> f2 env' b -> c)
-> f1 env a -> f2 env b -> Extend ArrayR OpenAcc env env' -> c
into2 f1 env' a -> f2 env' b -> c
op f1 env a
a f2 env b
b Extend ArrayR OpenAcc env env'
env = f1 env' a -> f2 env' b -> c
op (Extend ArrayR OpenAcc env env' -> f1 env a -> f1 env' a
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f1 env a
a) (Extend ArrayR OpenAcc env env' -> f2 env b -> f2 env' b
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc env env'
env f2 env b
b)


-- Replicate as a backwards permutation
--
-- TODO: If we have a pattern such as `replicate sh (map f xs)` then in some
--       cases it might be beneficial to not fuse these terms, if `f` is
--       expensive and/or `sh` is large.
--
replicateD
    :: HasCallStack
    => SliceIndex slix sl co sh
    -> Exp          aenv slix
    -> Cunctation r aenv (Array sl e)
    -> Cunctation D aenv (Array sh e)
replicateD :: forall slix sl co sh aenv r e.
HasCallStack =>
SliceIndex slix sl co sh
-> Exp aenv slix
-> Cunctation r aenv (Array sl e)
-> Cunctation D aenv (Array sh e)
replicateD SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix Cunctation r aenv (Array sl e)
cc
  = Text
-> Cunctation D aenv (Array sh e) -> Cunctation D aenv (Array sh e)
forall a. Text -> a -> a
Stats.ruleFired Text
"replicateD"
  (Cunctation D aenv (Array sh e) -> Cunctation D aenv (Array sh e))
-> Cunctation D aenv (Array sh e) -> Cunctation D aenv (Array sh e)
forall a b. (a -> b) -> a -> b
$ ShapeR sh
-> Exp aenv sh
-> Fun aenv (sh -> sl)
-> Cunctation r aenv (Array sl e)
-> Cunctation D aenv (Array sh e)
forall sh' aenv sh r e.
HasCallStack =>
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
backpermuteD (SliceIndex slix sl co sh -> ShapeR sh
forall slix sl co dim. SliceIndex slix sl co dim -> ShapeR dim
sliceDomainR SliceIndex slix sl co sh
sliceIndex) (SliceIndex slix sl co sh
-> Exp aenv slix -> OpenExp () aenv sl -> Exp aenv sh
forall slix sl co t env aenv.
SliceIndex slix sl co t
-> OpenExp env aenv slix
-> OpenExp env aenv sl
-> OpenExp env aenv t
IndexFull SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix (Cunctation r aenv (Array sl e) -> OpenExp () aenv sl
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Exp aenv sh
shape Cunctation r aenv (Array sl e)
cc)) (SliceIndex slix sl co sh -> Exp aenv slix -> Fun aenv (sh -> sl)
forall slix sl co sh aenv.
SliceIndex slix sl co sh -> Exp aenv slix -> Fun aenv (sh -> sl)
extend SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix) Cunctation r aenv (Array sl e)
cc


-- Dimensional slice as a backwards permutation
--
sliceD
    :: HasCallStack
    => SliceIndex slix sl co sh
    -> Exp          aenv slix
    -> Cunctation r aenv (Array sh e)
    -> Cunctation D aenv (Array sl e)
sliceD :: forall slix sl co sh aenv r e.
HasCallStack =>
SliceIndex slix sl co sh
-> Exp aenv slix
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sl e)
sliceD SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix Cunctation r aenv (Array sh e)
cc
  = Text
-> Cunctation D aenv (Array sl e) -> Cunctation D aenv (Array sl e)
forall a. Text -> a -> a
Stats.ruleFired Text
"sliceD"
  (Cunctation D aenv (Array sl e) -> Cunctation D aenv (Array sl e))
-> Cunctation D aenv (Array sl e) -> Cunctation D aenv (Array sl e)
forall a b. (a -> b) -> a -> b
$ ShapeR sl
-> Exp aenv sl
-> Fun aenv (sl -> sh)
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sl e)
forall sh' aenv sh r e.
HasCallStack =>
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
backpermuteD (SliceIndex slix sl co sh -> ShapeR sl
forall slix sl co dim. SliceIndex slix sl co dim -> ShapeR sl
sliceShapeR SliceIndex slix sl co sh
sliceIndex) (SliceIndex slix sl co sh
-> Exp aenv slix -> OpenExp () aenv sh -> Exp aenv sl
forall slix t co sh env aenv.
SliceIndex slix t co sh
-> OpenExp env aenv slix
-> OpenExp env aenv sh
-> OpenExp env aenv t
IndexSlice SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix (Cunctation r aenv (Array sh e) -> OpenExp () aenv sh
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Exp aenv sh
shape Cunctation r aenv (Array sh e)
cc)) (SliceIndex slix sl co sh -> Exp aenv slix -> Fun aenv (sl -> sh)
forall slix sl co sh aenv.
SliceIndex slix sl co sh -> Exp aenv slix -> Fun aenv (sl -> sh)
restrict SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix) Cunctation r aenv (Array sh e)
cc


-- Reshape an array
--
-- For delayed arrays this is implemented as an index space transformation. For
-- manifest arrays this can be done with the standard Reshape operation in
-- constant time without executing any array operations. This does not affect
-- the fusion process since the term is already manifest.
--
-- TLM: there was a runtime check to ensure the old and new shapes contained the
--      same number of elements: this has been lost for the delayed cases!
--
reshapeD
    :: HasCallStack
    => ShapeR sl
    -> Embed OpenAcc aenv (Array sh e)
    -> Exp           aenv sl
    -> Embed OpenAcc aenv (Array sl e)
reshapeD :: forall sl aenv sh e.
HasCallStack =>
ShapeR sl
-> Embed OpenAcc aenv (Array sh e)
-> Exp aenv sl
-> Embed OpenAcc aenv (Array sl e)
reshapeD ShapeR sl
slr (Embed Extend ArrayR OpenAcc aenv aenv'
env Cunctation r aenv' (Array sh e)
cc) (Extend ArrayR OpenAcc aenv aenv'
-> Exp aenv sl -> OpenExp () aenv' sl
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t.
Sink f =>
Extend s acc env env' -> f env t -> f env' t
sinkA Extend ArrayR OpenAcc aenv aenv'
env -> OpenExp () aenv' sl
sl)
  | Done ArrayVars aenv' (Array sh e)
v <- Cunctation r aenv' (Array sh e)
cc
  = Extend ArrayR OpenAcc aenv (aenv', Array sl e)
-> Cunctation M (aenv', Array sl e) (Array sl e)
-> Embed OpenAcc aenv (Array sl e)
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
env Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc aenv' (Array sl e)
-> Extend ArrayR OpenAcc aenv (aenv', Array sl e)
forall (acc :: * -> * -> *) aenv aenv' sh e.
HasArraysR acc =>
Extend ArrayR acc aenv aenv'
-> acc aenv' (Array sh e)
-> Extend ArrayR acc aenv (aenv', Array sh e)
`pushArrayEnv` PreOpenAcc OpenAcc aenv' (Array sl e) -> OpenAcc aenv' (Array sl e)
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (ShapeR sl
-> OpenExp () aenv' sl
-> OpenAcc aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sl e)
forall sh aenv (acc :: * -> * -> *) sh' e.
ShapeR sh
-> Exp aenv sh
-> acc aenv (Array sh' e)
-> PreOpenAcc acc aenv (Array sh e)
Reshape ShapeR sl
slr OpenExp () aenv' sl
sl ((forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t)
-> ArrayVars aenv' (Array sh e) -> OpenAcc aenv' (Array sh e)
forall (acc :: * -> * -> *) aenv arrs.
InjectAcc acc -> ArrayVars aenv arrs -> acc aenv arrs
avarsIn PreOpenAcc OpenAcc env t -> OpenAcc env t
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc ArrayVars aenv' (Array sh e)
v))) (Cunctation M (aenv', Array sl e) (Array sl e)
 -> Embed OpenAcc aenv (Array sl e))
-> Cunctation M (aenv', Array sl e) (Array sl e)
-> Embed OpenAcc aenv (Array sl e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sl e)
-> Cunctation M (aenv', Array sl e) (Array sl e)
forall sh e aenv.
ArrayR (Array sh e) -> Cunctation M (aenv, Array sh e) (Array sh e)
doneZeroIdx ArrayR (Array sl e)
repr

  | Bool
otherwise
  = Text
-> Embed OpenAcc aenv (Array sl e)
-> Embed OpenAcc aenv (Array sl e)
forall a. Text -> a -> a
Stats.ruleFired Text
"reshapeD"
  (Embed OpenAcc aenv (Array sl e)
 -> Embed OpenAcc aenv (Array sl e))
-> Embed OpenAcc aenv (Array sl e)
-> Embed OpenAcc aenv (Array sl e)
forall a b. (a -> b) -> a -> b
$ Extend ArrayR OpenAcc aenv aenv'
-> Cunctation D aenv' (Array sl e)
-> Embed OpenAcc aenv (Array sl e)
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env (ShapeR sl
-> OpenExp () aenv' sl
-> Fun aenv' (sl -> sh)
-> Cunctation r aenv' (Array sh e)
-> Cunctation D aenv' (Array sl e)
forall sh' aenv sh r e.
HasCallStack =>
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Cunctation r aenv (Array sh e)
-> Cunctation D aenv (Array sh' e)
backpermuteD ShapeR sl
slr OpenExp () aenv' sl
sl (ShapeR sh
-> OpenExp () aenv' sh
-> ShapeR sl
-> OpenExp () aenv' sl
-> Fun aenv' (sl -> sh)
forall sh' env aenv sh.
ShapeR sh'
-> OpenExp env aenv sh'
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenFun env aenv (sh -> sh')
reindex (ArrayR (Array sh e) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape (ArrayR (Array sh e) -> ShapeR sh)
-> ArrayR (Array sh e) -> ShapeR sh
forall a b. (a -> b) -> a -> b
$ Cunctation r aenv' (Array sh e) -> ArrayR (Array sh e)
forall (f :: * -> * -> *) aenv sh e.
HasArraysR f =>
f aenv (Array sh e) -> ArrayR (Array sh e)
arrayR Cunctation r aenv' (Array sh e)
cc) (Cunctation r aenv' (Array sh e) -> OpenExp () aenv' sh
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Exp aenv sh
shape Cunctation r aenv' (Array sh e)
cc) ShapeR sl
slr OpenExp () aenv' sl
sl) Cunctation r aenv' (Array sh e)
cc)

  where
    ArrayR ShapeR sh
_ TypeR e
TypeR e
tR = Cunctation r aenv' (Array sh e) -> ArrayR (Array sh e)
forall (f :: * -> * -> *) aenv sh e.
HasArraysR f =>
f aenv (Array sh e) -> ArrayR (Array sh e)
arrayR Cunctation r aenv' (Array sh e)
cc
    repr :: ArrayR (Array sl e)
repr        = ShapeR sl -> TypeR e -> ArrayR (Array sl e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sl
slr TypeR e
tR


-- Combine two arrays element-wise with a binary function to produce a delayed
-- array.
--
zipWithD
    :: HasCallStack
    => TypeR c
    -> Fun          aenv (a -> b -> c)
    -> Cunctation r aenv (Array sh a)
    -> Cunctation s aenv (Array sh b)
    -> Cunctation D aenv (Array sh c)
zipWithD :: forall c aenv a b r sh s.
HasCallStack =>
TypeR c
-> Fun aenv (a -> b -> c)
-> Cunctation r aenv (Array sh a)
-> Cunctation s aenv (Array sh b)
-> Cunctation D aenv (Array sh c)
zipWithD TypeR c
tR Fun aenv (a -> b -> c)
f Cunctation r aenv (Array sh a)
cc1 Cunctation s aenv (Array sh b)
cc0
  -- Two stepper functions identically accessing the same array can be kept in
  -- stepping form. This might yield a simpler final term.
  --
  | Step (ArrayR ShapeR sh
shR TypeR e
_) Exp aenv sh'
sh1 Fun aenv (sh' -> sh)
p1 Fun aenv (a -> b)
f1 ArrayVar aenv (Array sh a)
v1 <- Cunctation r aenv (Array sh a) -> Cunctation D aenv (Array sh a)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
delaying Cunctation r aenv (Array sh a)
cc1
  , Step ArrayR (Array sh' b)
_              Exp aenv sh'
sh0 Fun aenv (sh' -> sh)
p0 Fun aenv (a -> b)
f0 ArrayVar aenv (Array sh a)
v0 <- Cunctation s aenv (Array sh b) -> Cunctation D aenv (Array sh b)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
delaying Cunctation s aenv (Array sh b)
cc0
  , Just Array sh a :~: Array sh a
Refl                        <- ArrayVar aenv (Array sh a)
-> ArrayVar aenv (Array sh a) -> Maybe (Array sh a :~: Array sh a)
forall (s :: * -> *) env t1 t2.
Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2)
matchVar ArrayVar aenv (Array sh a)
v1 ArrayVar aenv (Array sh a)
v0
  , Just (sh' -> sh) :~: (sh' -> sh)
Refl                        <- Fun aenv (sh' -> sh)
-> Fun aenv (sh' -> sh) -> Maybe ((sh' -> sh) :~: (sh' -> sh))
forall env aenv s t.
OpenFun env aenv s -> OpenFun env aenv t -> Maybe (s :~: t)
matchOpenFun Fun aenv (sh' -> sh)
p1 Fun aenv (sh' -> sh)
p0
  = Text
-> Cunctation D aenv (Array sh c) -> Cunctation D aenv (Array sh c)
forall a. Text -> a -> a
Stats.ruleFired Text
"zipWithD/step"
  (Cunctation D aenv (Array sh c) -> Cunctation D aenv (Array sh c))
-> Cunctation D aenv (Array sh c) -> Cunctation D aenv (Array sh c)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh c)
-> Exp aenv sh
-> Fun aenv (sh -> sh)
-> Fun aenv (a -> c)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array sh c)
forall env' e aenv sh a.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> sh)
-> Fun aenv (a -> e)
-> ArrayVar aenv (Array sh a)
-> Cunctation D aenv (Array env' e)
Step (ShapeR sh -> TypeR c -> ArrayR (Array sh c)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh
shR TypeR c
tR) (ShapeR sh -> Exp aenv sh -> Exp aenv sh -> Exp aenv sh
forall sh env aenv.
ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
intersect ShapeR sh
ShapeR sh
shR Exp aenv sh
Exp aenv sh'
sh1 Exp aenv sh
Exp aenv sh'
sh0) Fun aenv (sh -> sh)
Fun aenv (sh' -> sh)
p0 (Fun aenv (a -> b -> c)
-> Fun aenv (a -> a) -> Fun aenv (a -> b) -> Fun aenv (a -> c)
forall aenv a b c e.
HasCallStack =>
Fun aenv (a -> b -> c)
-> Fun aenv (e -> a) -> Fun aenv (e -> b) -> Fun aenv (e -> c)
combine Fun aenv (a -> b -> c)
f Fun aenv (a -> a)
Fun aenv (a -> b)
f1 Fun aenv (a -> b)
Fun aenv (a -> b)
f0) ArrayVar aenv (Array sh a)
ArrayVar aenv (Array sh a)
v0

  -- Otherwise transform both delayed terms into (index -> value) mappings and
  -- combine the two indexing functions that way.
  --
  | Yield (ArrayR ShapeR sh
shR TypeR e
_) Exp aenv sh
sh1 Fun aenv (sh -> e)
f1 <- Cunctation r aenv (Array sh a) -> Cunctation D aenv (Array sh a)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
yield Cunctation r aenv (Array sh a)
cc1
  , Yield ArrayR (Array sh e)
_              Exp aenv sh
sh0 Fun aenv (sh -> e)
f0 <- Cunctation s aenv (Array sh b) -> Cunctation D aenv (Array sh b)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
yield Cunctation s aenv (Array sh b)
cc0
  = Text
-> Cunctation D aenv (Array sh c) -> Cunctation D aenv (Array sh c)
forall a. Text -> a -> a
Stats.ruleFired Text
"zipWithD"
  (Cunctation D aenv (Array sh c) -> Cunctation D aenv (Array sh c))
-> Cunctation D aenv (Array sh c) -> Cunctation D aenv (Array sh c)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh c)
-> Exp aenv sh
-> Fun aenv (sh -> c)
-> Cunctation D aenv (Array sh c)
forall env' e aenv.
ArrayR (Array env' e)
-> Exp aenv env'
-> Fun aenv (env' -> e)
-> Cunctation D aenv (Array env' e)
Yield (ShapeR sh -> TypeR c -> ArrayR (Array sh c)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh
shR TypeR c
tR) (ShapeR sh -> Exp aenv sh -> Exp aenv sh -> Exp aenv sh
forall sh env aenv.
ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
intersect ShapeR sh
ShapeR sh
shR Exp aenv sh
Exp aenv sh
sh1 Exp aenv sh
Exp aenv sh
sh0) (Fun aenv (a -> b -> c)
-> Fun aenv (sh -> a) -> Fun aenv (sh -> b) -> Fun aenv (sh -> c)
forall aenv a b c e.
HasCallStack =>
Fun aenv (a -> b -> c)
-> Fun aenv (e -> a) -> Fun aenv (e -> b) -> Fun aenv (e -> c)
combine Fun aenv (a -> b -> c)
f Fun aenv (sh -> a)
Fun aenv (sh -> e)
f1 Fun aenv (sh -> b)
Fun aenv (sh -> e)
f0)

  | Bool
otherwise
  = [Char] -> Cunctation D aenv (Array sh c)
forall a. HasCallStack => [Char] -> a
error [Char]
"work is stressing me out, I should take a break"
  where
    combine :: forall aenv a b c e. HasCallStack
            => Fun aenv (a -> b -> c)
            -> Fun aenv (e -> a)
            -> Fun aenv (e -> b)
            -> Fun aenv (e -> c)
    combine :: forall aenv a b c e.
HasCallStack =>
Fun aenv (a -> b -> c)
-> Fun aenv (e -> a) -> Fun aenv (e -> b) -> Fun aenv (e -> c)
combine Fun aenv (a -> b -> c)
c Fun aenv (e -> a)
ixa Fun aenv (e -> b)
ixb
      | Lam ELeftHandSide a () env'
lhs1 (Body OpenExp env' aenv t1
ixa') <- Fun aenv (e -> a)
ixa
      , Lam ELeftHandSide a () env'
lhs2 (Body OpenExp env' aenv t1
ixb') <- Fun aenv (e -> b)
ixb
      -- The two LeftHandSides may differ in the use of wildcards. If they do not match, we must
      -- combine them as done in `combineLHS`. As this will probably not occur often and requires
      -- additional weakening, we do a quick check whether the left hand sides are equal.
      --
      = case ELeftHandSide a () env'
-> ELeftHandSide a () env'
-> Maybe (ELeftHandSide a () env' :~: ELeftHandSide a () env')
forall env env1 env2 t1 t2.
ELeftHandSide t1 env env1
-> ELeftHandSide t2 env env2
-> Maybe (ELeftHandSide t1 env env1 :~: ELeftHandSide t2 env env2)
matchELeftHandSide ELeftHandSide a () env'
lhs1 ELeftHandSide a () env'
lhs2 of
          Just ELeftHandSide a () env' :~: ELeftHandSide a () env'
Refl
            | Lam ELeftHandSide a env' env'
lhsA (Lam ELeftHandSide a env' env'
lhsB (Body OpenExp env' aenv t1
c')) <- (() :> env')
-> Fun aenv (a -> b -> c) -> OpenFun env' aenv (a -> b -> c)
forall env env' aenv t.
(env :> env') -> OpenFun env aenv t -> OpenFun env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE (ELeftHandSide a () env' -> () :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ELeftHandSide a () env'
lhs1) Fun aenv (a -> b -> c)
c
              -> ELeftHandSide e () env' -> OpenFun env' aenv c -> Fun aenv (e -> c)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam ELeftHandSide e () env'
ELeftHandSide a () env'
lhs1 (OpenFun env' aenv c -> Fun aenv (e -> c))
-> OpenFun env' aenv c -> Fun aenv (e -> c)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv c -> OpenFun env' aenv c
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv c -> OpenFun env' aenv c)
-> OpenExp env' aenv c -> OpenFun env' aenv c
forall a b. (a -> b) -> a -> b
$ ELeftHandSide a env' env'
-> OpenExp env' aenv a
-> OpenExp env' aenv c
-> OpenExp env' aenv c
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide a env' env'
lhsA OpenExp env' aenv t1
OpenExp env' aenv a
ixa'  (OpenExp env' aenv c -> OpenExp env' aenv c)
-> OpenExp env' aenv c -> OpenExp env' aenv c
forall a b. (a -> b) -> a -> b
$ ELeftHandSide a env' env'
-> OpenExp env' aenv a
-> OpenExp env' aenv c
-> OpenExp env' aenv c
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide a env' env'
lhsB ((env' :> env') -> OpenExp env' aenv a -> OpenExp env' aenv a
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE (ELeftHandSide a env' env' -> env' :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ELeftHandSide a env' env'
lhsA)       OpenExp env' aenv a
OpenExp env' aenv t1
ixb') OpenExp env' aenv c
OpenExp env' aenv t1
c'
          Maybe (ELeftHandSide a () env' :~: ELeftHandSide a () env')
Nothing
            | CombinedLHS LeftHandSide ScalarType a () env'
lhs env' :> env'
k1 env' :> env'
k2         <- ELeftHandSide a () env'
-> LeftHandSide ScalarType a () env'
-> CombinedLHS ScalarType a env' env' ()
forall (s :: * -> *) t env env1' env2'.
HasCallStack =>
LeftHandSide s t env env1'
-> LeftHandSide s t env env2' -> CombinedLHS s t env1' env2' env
combineLHS ELeftHandSide a () env'
lhs1 LeftHandSide ScalarType a () env'
ELeftHandSide a () env'
lhs2
            , Lam ELeftHandSide a env' env'
lhsA (Lam ELeftHandSide a env' env'
lhsB (Body OpenExp env' aenv t1
c')) <- (() :> env')
-> Fun aenv (a -> b -> c) -> OpenFun env' aenv (a -> b -> c)
forall env env' aenv t.
(env :> env') -> OpenFun env aenv t -> OpenFun env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE (LeftHandSide ScalarType a () env' -> () :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS LeftHandSide ScalarType a () env'
lhs) Fun aenv (a -> b -> c)
c
            , OpenExp env' aenv t1
ixa''                         <- (env' :> env') -> OpenExp env' aenv t1 -> OpenExp env' aenv t1
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env' :> env'
k1 OpenExp env' aenv t1
ixa'
              -> ELeftHandSide e () env' -> OpenFun env' aenv c -> Fun aenv (e -> c)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam ELeftHandSide e () env'
LeftHandSide ScalarType a () env'
lhs  (OpenFun env' aenv c -> Fun aenv (e -> c))
-> OpenFun env' aenv c -> Fun aenv (e -> c)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv c -> OpenFun env' aenv c
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv c -> OpenFun env' aenv c)
-> OpenExp env' aenv c -> OpenFun env' aenv c
forall a b. (a -> b) -> a -> b
$ ELeftHandSide a env' env'
-> OpenExp env' aenv a
-> OpenExp env' aenv c
-> OpenExp env' aenv c
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide a env' env'
lhsA OpenExp env' aenv t1
OpenExp env' aenv a
ixa'' (OpenExp env' aenv c -> OpenExp env' aenv c)
-> OpenExp env' aenv c -> OpenExp env' aenv c
forall a b. (a -> b) -> a -> b
$ ELeftHandSide a env' env'
-> OpenExp env' aenv a
-> OpenExp env' aenv c
-> OpenExp env' aenv c
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide a env' env'
lhsB ((env' :> env') -> OpenExp env' aenv a -> OpenExp env' aenv a
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE (ELeftHandSide a env' env' -> env' :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ELeftHandSide a env' env'
lhsA (env' :> env') -> (env' :> env') -> env' :> env'
forall env2 env3 env1.
(env2 :> env3) -> (env1 :> env2) -> env1 :> env3
.> env' :> env'
k2) OpenExp env' aenv t1
OpenExp env' aenv a
ixb') OpenExp env' aenv c
OpenExp env' aenv t1
c'
          Maybe (ELeftHandSide a () env' :~: ELeftHandSide a () env')
_
              -> [Char] -> Fun aenv (e -> c)
forall a. HasCallStack => [Char] -> a
error [Char]
"how's your break?"
      --
      | Bool
otherwise
      = [Char] -> Fun aenv (e -> c)
forall a. HasCallStack => [Char] -> a
error [Char]
"work is stressing me out, I should get back to it"


data CombinedLHS s t env1' env2' env where
  CombinedLHS :: LeftHandSide s t env env'
              -> env1' :> env'
              -> env2' :> env'
              -> CombinedLHS s t env1' env2' env

combineLHS
    :: HasCallStack
    => LeftHandSide s t env env1'
    -> LeftHandSide s t env env2'
    -> CombinedLHS  s t env1' env2' env
combineLHS :: forall (s :: * -> *) t env env1' env2'.
HasCallStack =>
LeftHandSide s t env env1'
-> LeftHandSide s t env env2' -> CombinedLHS s t env1' env2' env
combineLHS = (env :> env)
-> (env :> env)
-> LeftHandSide s t env env1'
-> LeftHandSide s t env env2'
-> CombinedLHS s t env1' env2' env
forall env1 env env2 (s :: * -> *) t env1' env2'.
(env1 :> env)
-> (env2 :> env)
-> LeftHandSide s t env1 env1'
-> LeftHandSide s t env2 env2'
-> CombinedLHS s t env1' env2' env
go env :> env
forall env. env :> env
weakenId env :> env
forall env. env :> env
weakenId
  where
    go :: env1 :> env -> env2 :> env -> LeftHandSide s t env1 env1' -> LeftHandSide s t env2 env2' -> CombinedLHS s t env1' env2' env
    go :: forall env1 env env2 (s :: * -> *) t env1' env2'.
(env1 :> env)
-> (env2 :> env)
-> LeftHandSide s t env1 env1'
-> LeftHandSide s t env2 env2'
-> CombinedLHS s t env1' env2' env
go env1 :> env
k1 env2 :> env
k2 (LeftHandSideWildcard TupR s t
tR) (LeftHandSideWildcard TupR s t
_) = LeftHandSide s t env env
-> (env1' :> env)
-> (env2' :> env)
-> CombinedLHS s t env1' env2' env
forall (s :: * -> *) t env env' env1' env2'.
LeftHandSide s t env env'
-> (env1' :> env')
-> (env2' :> env')
-> CombinedLHS s t env1' env2' env
CombinedLHS (TupR s t -> LeftHandSide s t env env
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard TupR s t
tR)    env1 :> env
env1' :> env
k1        env2 :> env
env2' :> env
k2
    go env1 :> env
k1 env2 :> env
k2 (LeftHandSideSingle s t
tR)   (LeftHandSideSingle s t
_)   = LeftHandSide s t env (env, t)
-> (env1' :> (env, t))
-> (env2' :> (env, t))
-> CombinedLHS s t env1' env2' env
forall (s :: * -> *) t env env' env1' env2'.
LeftHandSide s t env env'
-> (env1' :> env')
-> (env2' :> env')
-> CombinedLHS s t env1' env2' env
CombinedLHS (s t -> LeftHandSide s t env (env, t)
forall (s :: * -> *) v env. s v -> LeftHandSide s v env (env, v)
LeftHandSideSingle s t
tR)      ((env1 :> env) -> (env1, t) :> (env, t)
forall env env' t. (env :> env') -> (env, t) :> (env', t)
sink env1 :> env
k1) ((env2 :> env) -> (env2, t) :> (env, t)
forall env env' t. (env :> env') -> (env, t) :> (env', t)
sink env2 :> env
k2)
    go env1 :> env
k1 env2 :> env
k2 (LeftHandSidePair LeftHandSide s v1 env1 env'1
l1 LeftHandSide s v2 env'1 env1'
h1)  (LeftHandSidePair LeftHandSide s v1 env2 env'1
l2 LeftHandSide s v2 env'1 env2'
h2)
      | CombinedLHS LeftHandSide s v1 env env'
l env'1 :> env'
k1'  env'1 :> env'
k2'  <- (env1 :> env)
-> (env2 :> env)
-> LeftHandSide s v1 env1 env'1
-> LeftHandSide s v1 env2 env'1
-> CombinedLHS s v1 env'1 env'1 env
forall env1 env env2 (s :: * -> *) t env1' env2'.
(env1 :> env)
-> (env2 :> env)
-> LeftHandSide s t env1 env1'
-> LeftHandSide s t env2 env2'
-> CombinedLHS s t env1' env2' env
go env1 :> env
k1  env2 :> env
k2  LeftHandSide s v1 env1 env'1
l1 LeftHandSide s v1 env2 env'1
LeftHandSide s v1 env2 env'1
l2
      , CombinedLHS LeftHandSide s v2 env' env'
h env1' :> env'
k1'' env2' :> env'
k2'' <- (env'1 :> env')
-> (env'1 :> env')
-> LeftHandSide s v2 env'1 env1'
-> LeftHandSide s v2 env'1 env2'
-> CombinedLHS s v2 env1' env2' env'
forall env1 env env2 (s :: * -> *) t env1' env2'.
(env1 :> env)
-> (env2 :> env)
-> LeftHandSide s t env1 env1'
-> LeftHandSide s t env2 env2'
-> CombinedLHS s t env1' env2' env
go env'1 :> env'
k1' env'1 :> env'
k2' LeftHandSide s v2 env'1 env1'
h1 LeftHandSide s v2 env'1 env2'
LeftHandSide s v2 env'1 env2'
h2             = LeftHandSide s t env env'
-> (env1' :> env')
-> (env2' :> env')
-> CombinedLHS s t env1' env2' env
forall (s :: * -> *) t env env' env1' env2'.
LeftHandSide s t env env'
-> (env1' :> env')
-> (env2' :> env')
-> CombinedLHS s t env1' env2' env
CombinedLHS (LeftHandSide s v1 env env'
-> LeftHandSide s v2 env' env' -> LeftHandSide s (v1, v2) env env'
forall (s :: * -> *) v1 env env'1 v2 env'.
LeftHandSide s v1 env env'1
-> LeftHandSide s v2 env'1 env' -> LeftHandSide s (v1, v2) env env'
LeftHandSidePair LeftHandSide s v1 env env'
l LeftHandSide s v2 env' env'
h)       env1' :> env'
k1''      env2' :> env'
k2''
    go env1 :> env
k1 env2 :> env
k2 (LeftHandSideWildcard TupR s t
_)  LeftHandSide s t env2 env2'
lhs
      | Exists LeftHandSide s t env a
lhs' <- LeftHandSide s t env2 env2' -> Exists (LeftHandSide s t env)
forall (s :: * -> *) t aenv1 aenv1' aenv2.
LeftHandSide s t aenv1 aenv1' -> Exists (LeftHandSide s t aenv2)
rebuildLHS LeftHandSide s t env2 env2'
lhs                           = LeftHandSide s t env a
-> (env1' :> a) -> (env2' :> a) -> CombinedLHS s t env1' env2' env
forall (s :: * -> *) t env env' env1' env2'.
LeftHandSide s t env env'
-> (env1' :> env')
-> (env2' :> env')
-> CombinedLHS s t env1' env2' env
CombinedLHS LeftHandSide s t env a
lhs'        (LeftHandSide s t env a -> env :> a
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS LeftHandSide s t env a
lhs' (env :> a) -> (env1' :> env) -> env1' :> a
forall env2 env3 env1.
(env2 :> env3) -> (env1 :> env2) -> env1 :> env3
.> env1 :> env
env1' :> env
k1) (LeftHandSide s t env2 env2'
-> LeftHandSide s t env a -> (env2 :> env) -> env2' :> a
forall (s :: * -> *) t env1 env1' env2 env2'.
HasCallStack =>
LeftHandSide s t env1 env1'
-> LeftHandSide s t env2 env2' -> (env1 :> env2) -> env1' :> env2'
sinkWithLHS LeftHandSide s t env2 env2'
lhs LeftHandSide s t env a
lhs' env2 :> env
k2)
    go env1 :> env
k1 env2 :> env
k2 LeftHandSide s t env1 env1'
lhs                       (LeftHandSideWildcard TupR s t
_)
      | Exists LeftHandSide s t env a
lhs' <- LeftHandSide s t env1 env1' -> Exists (LeftHandSide s t env)
forall (s :: * -> *) t aenv1 aenv1' aenv2.
LeftHandSide s t aenv1 aenv1' -> Exists (LeftHandSide s t aenv2)
rebuildLHS LeftHandSide s t env1 env1'
lhs                           = LeftHandSide s t env a
-> (env1' :> a) -> (env2' :> a) -> CombinedLHS s t env1' env2' env
forall (s :: * -> *) t env env' env1' env2'.
LeftHandSide s t env env'
-> (env1' :> env')
-> (env2' :> env')
-> CombinedLHS s t env1' env2' env
CombinedLHS LeftHandSide s t env a
lhs'        (LeftHandSide s t env1 env1'
-> LeftHandSide s t env a -> (env1 :> env) -> env1' :> a
forall (s :: * -> *) t env1 env1' env2 env2'.
HasCallStack =>
LeftHandSide s t env1 env1'
-> LeftHandSide s t env2 env2' -> (env1 :> env2) -> env1' :> env2'
sinkWithLHS LeftHandSide s t env1 env1'
lhs LeftHandSide s t env a
lhs' env1 :> env
k1)  (LeftHandSide s t env a -> env :> a
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS LeftHandSide s t env a
lhs' (env :> a) -> (env2' :> env) -> env2' :> a
forall env2 env3 env1.
(env2 :> env3) -> (env1 :> env2) -> env1 :> env3
.> env2 :> env
env2' :> env
k2)
    go env1 :> env
_ env2 :> env
_ LeftHandSide s t env1 env1'
_ LeftHandSide s t env2 env2'
_
      = Format
  (CombinedLHS s t env1' env2' env) (CombinedLHS s t env1' env2' env)
-> CombinedLHS s t env1' env2' env
forall r a. HasCallStack => Format r a -> a
internalError Format
  (CombinedLHS s t env1' env2' env) (CombinedLHS s t env1' env2' env)
"unexpected LHS combination"



-- NOTE: [Sharing vs. Fusion]
--
-- The approach to array fusion is similar to that the first generation of Repa.
-- It was discovered that the most immediately pressing problem with delayed
-- arrays in Repa-1 was that it did not preserve sharing of collective
-- operations, leading to excessive recomputation and severe repercussions on
-- performance if the user did not explicitly intervene.
--
-- However, as we have explicit sharing information in the term tree, so it is
-- straightforward to respect sharing by not fusing let-bindings, as that
-- introduces work duplication. However, sometimes we can be cleverer.
--
-- let-floating:
-- -------------
--
-- If the binding is of manifest data, we can instead move the let-binding to a
-- different point in the program and then continue to fuse into the body. This
-- is done by adding the bound term to the Extend environment. In essence this
-- is covering a different occurrence of the same problem Extend was introduced
-- to handle: let bindings of manifest data unnecessarily get in the way of the
-- fusion process. For example:
--
--   map f (zipWith g xs (map h xs))
--
-- after sharing recovery results in:
--
--   map f (let a0 = xs in zipWith g a0 (map h a0))
--
-- Without allowing the binding for a0 to float outwards, `map f` will not be
-- fused into the rest of the program.
--
-- let-elimination:
-- ----------------
--
-- Array binding points appear in the program because the array data _or_ shape
-- was accessed multiple times in the source program. In general we want to fuse
-- arbitrary sequences of array _data_, irrespective of how the shape component
-- is used. For example, reverse is defined in the prelude as:
--
--   reverse xs = let len  = unindex1 (shape xs)
--                    pf i = len - i - 1
--                in
--                backpermute (shape xs) (ilift1 pf) xs
--
-- Sharing recovery introduces a let-binding for the input `xs` since it is used
-- thrice in the definition, which impedes subsequent fusion. However the actual
-- array data is only accessed once, with the remaining two uses querying the
-- array shape. Since the delayed terms contain the shape of the array they
-- represent as a scalar term, if the data component otherwise satisfies the
-- rules for fusing terms, as it does in this example, we can eliminate the
-- let-binding by pushing the scalar shape and value generation terms directly
-- into the body.
--
-- Let-elimination can also be used to _introduce_ work duplication, which may
-- be beneficial if we can estimate that the cost of re-computation is less than
-- the cost of completely evaluating the array and subsequently retrieving the
-- data from memory.
--
-- let-binding:
-- ------------
--
-- Ultimately, we might not want to eliminate the binding. If so, evaluate it
-- and add it to a _clean_ Extend environment for the body. If not, the Extend
-- list effectively _flattens_ all bindings, so any terms required for the bound
-- term get lifted out to the same scope as the body. This increases their
-- lifetime and hence raises the maximum memory used. If we don't do this, we
-- get terms such as:
--
--   let a0  = <terms for binding> in
--   let bnd = <bound term> in
--   <body term>
--
-- rather than the following, where the scope of a0 is clearly only availably
-- when evaluating the bound term, as it should be:
--
--   let bnd =
--     let a0 = <terms for binding>
--     in <bound term>
--   in <body term>
--
aletD :: HasCallStack
      => EmbedAcc OpenAcc
      -> ElimAcc  OpenAcc
      -> ALeftHandSide arrs aenv aenv'
      ->       OpenAcc aenv  arrs
      ->       OpenAcc aenv' brrs
      -> Embed OpenAcc aenv  brrs
aletD :: forall arrs aenv aenv' brrs.
HasCallStack =>
EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> ALeftHandSide arrs aenv aenv'
-> OpenAcc aenv arrs
-> OpenAcc aenv' brrs
-> Embed OpenAcc aenv brrs
aletD EmbedAcc OpenAcc
embedAcc ElimAcc OpenAcc
elimAcc ALeftHandSide arrs aenv aenv'
lhs (OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc -> Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
cc1) OpenAcc aenv' brrs
acc0

  -- let-floating
  -- ------------
  --
  -- Immediately inline the variable referring to the bound expression into the
  -- body, instead of adding to the environments and creating an indirection
  -- that must be later eliminated by shrinking.
  --
  | LeftHandSideSingle ArrayR arrs
_                  <- ALeftHandSide arrs aenv aenv'
lhs
  , Done (TupRsingle v1 :: Var ArrayR aenv' arrs
v1@(Var ArrayR{} Idx aenv' arrs
_)) <- Cunctation r aenv' arrs
cc1
  , Embed Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' brrs
cc0                        <- OpenAcc aenv' brrs -> Embed OpenAcc aenv' brrs
EmbedAcc OpenAcc
embedAcc (OpenAcc aenv' brrs -> Embed OpenAcc aenv' brrs)
-> OpenAcc aenv' brrs -> Embed OpenAcc aenv' brrs
forall a b. (a -> b) -> a -> b
$ (forall sh e.
 ArrayVar aenv' (Array sh e)
 -> PreOpenAcc (AccClo OpenAcc) aenv' (Array sh e))
-> OpenAcc aenv' brrs -> OpenAcc aenv' brrs
forall (f :: * -> * -> *) (fa :: (* -> * -> *) -> * -> * -> *) aenv
       aenv' a.
(Rebuildable f, SyntacticAcc fa) =>
(forall sh e.
 ArrayVar aenv (Array sh e) -> fa (AccClo f) aenv' (Array sh e))
-> f aenv a -> f aenv' a
forall (fa :: (* -> * -> *) -> * -> * -> *) aenv aenv' a.
SyntacticAcc fa =>
(forall sh e.
 ArrayVar aenv (Array sh e)
 -> fa (AccClo OpenAcc) aenv' (Array sh e))
-> OpenAcc aenv a -> OpenAcc aenv' a
rebuildA (PreOpenAcc OpenAcc aenv' (Array sh e)
-> ArrayVar (aenv', Array sh e) (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sh e)
forall (acc :: * -> * -> *) aenv t sh2 e2.
PreOpenAcc acc aenv t
-> ArrayVar (aenv, t) (Array sh2 e2)
-> PreOpenAcc acc aenv (Array sh2 e2)
subAtop (ArrayVar aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar Var ArrayR aenv' arrs
ArrayVar aenv' (Array sh e)
v1) (ArrayVar (aenv', Array sh e) (Array sh e)
 -> PreOpenAcc OpenAcc aenv' (Array sh e))
-> (ArrayVar aenv' (Array sh e)
    -> ArrayVar (aenv', Array sh e) (Array sh e))
-> ArrayVar aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sh e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Extend ArrayR OpenAcc aenv aenv'
-> Var ArrayR (aenv, Array sh e) (Array sh e)
-> ArrayVar (aenv', Array sh e) (Array sh e)
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t' t.
Sink f =>
Extend s acc env env' -> f (env, t') t -> f (env', t') t
sink1 Extend ArrayR OpenAcc aenv aenv'
env1) OpenAcc aenv' brrs
acc0
  = Text -> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a. Text -> a -> a
Stats.ruleFired Text
"aletD/float"
  (Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs)
-> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a b. (a -> b) -> a -> b
$ Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' brrs -> Embed OpenAcc aenv brrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
env1 Extend ArrayR OpenAcc aenv aenv'
-> Extend ArrayR OpenAcc aenv' aenv'
-> Extend ArrayR OpenAcc aenv aenv'
forall (s :: * -> *) (acc :: * -> * -> *) env env' env''.
Extend s acc env env'
-> Extend s acc env' env'' -> Extend s acc env env''
`append` Extend ArrayR OpenAcc aenv' aenv'
env0) Cunctation r aenv' brrs
cc0

  -- Ensure we only call 'embedAcc' once on the body expression
  --
  | Bool
otherwise
  = EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> ALeftHandSide arrs aenv aenv'
-> Embed OpenAcc aenv arrs
-> Embed OpenAcc aenv' brrs
-> Embed OpenAcc aenv brrs
forall aenv aenv' arrs brrs.
HasCallStack =>
EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> ALeftHandSide arrs aenv aenv'
-> Embed OpenAcc aenv arrs
-> Embed OpenAcc aenv' brrs
-> Embed OpenAcc aenv brrs
aletD' OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv s -> OpenAcc (aenv, s) t -> Bool
ElimAcc OpenAcc
elimAcc ALeftHandSide arrs aenv aenv'
lhs (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' arrs -> Embed OpenAcc aenv arrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
cc1) (OpenAcc aenv' brrs -> Embed OpenAcc aenv' brrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv' brrs
acc0)


aletD' :: forall aenv aenv' arrs brrs. HasCallStack
       => EmbedAcc OpenAcc
       -> ElimAcc OpenAcc
       -> ALeftHandSide arrs aenv aenv'
       -> Embed OpenAcc aenv  arrs
       -> Embed OpenAcc aenv' brrs
       -> Embed OpenAcc aenv  brrs
aletD' :: forall aenv aenv' arrs brrs.
HasCallStack =>
EmbedAcc OpenAcc
-> ElimAcc OpenAcc
-> ALeftHandSide arrs aenv aenv'
-> Embed OpenAcc aenv arrs
-> Embed OpenAcc aenv' brrs
-> Embed OpenAcc aenv brrs
aletD' EmbedAcc OpenAcc
embedAcc ElimAcc OpenAcc
elimAcc (LeftHandSideSingle ArrayR{}) (Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
cc1) (Embed Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' brrs
cc0)

  -- let-binding
  -- -----------
  --
  -- Check whether we can eliminate the let-binding. Note that we must inspect
  -- the entire term, not just the Cunctation that would be produced by
  -- embedAcc. If we don't we can be left with dead terms that don't get
  -- eliminated. This problem occurred in the canny program.
  --
  | OpenAcc aenv arrs
acc1    <- Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' arrs -> Embed OpenAcc aenv arrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
cc1)
  , Bool
False   <- OpenAcc aenv arrs -> OpenAcc (aenv, arrs) brrs -> Bool
ElimAcc OpenAcc
elimAcc OpenAcc aenv arrs
acc1 OpenAcc aenv' brrs
OpenAcc (aenv, arrs) brrs
acc0
  = Text -> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a. Text -> a -> a
Stats.ruleFired Text
"aletD/bind"
  (Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs)
-> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a b. (a -> b) -> a -> b
$ Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' brrs -> Embed OpenAcc aenv brrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv
forall (s :: * -> *) (f :: * -> * -> *) env. Extend s f env env
BaseEnv Extend ArrayR OpenAcc aenv aenv
-> OpenAcc aenv (Array sh e)
-> Extend ArrayR OpenAcc aenv (aenv, Array sh e)
forall (acc :: * -> * -> *) aenv aenv' sh e.
HasArraysR acc =>
Extend ArrayR acc aenv aenv'
-> acc aenv' (Array sh e)
-> Extend ArrayR acc aenv (aenv', Array sh e)
`pushArrayEnv` OpenAcc aenv arrs
OpenAcc aenv (Array sh e)
acc1 Extend ArrayR OpenAcc aenv (aenv, Array sh e)
-> Extend ArrayR OpenAcc (aenv, Array sh e) aenv'
-> Extend ArrayR OpenAcc aenv aenv'
forall (s :: * -> *) (acc :: * -> * -> *) env env' env''.
Extend s acc env env'
-> Extend s acc env' env'' -> Extend s acc env env''
`append` Extend ArrayR OpenAcc aenv' aenv'
Extend ArrayR OpenAcc (aenv, Array sh e) aenv'
env0) Cunctation r aenv' brrs
cc0

  -- let-elimination
  -- ---------------
  --
  -- Handle the remaining cases in a separate function. It turns out that this
  -- is important so we aren't excessively sinking/delaying terms.
  --
  | OpenAcc (aenv', Array sh e) brrs
acc0'   <- Extend ArrayR OpenAcc aenv aenv'
-> OpenAcc (aenv, Array sh e) brrs
-> OpenAcc (aenv', Array sh e) brrs
forall (f :: * -> * -> *) (s :: * -> *) (acc :: * -> * -> *) env
       env' t' t.
Sink f =>
Extend s acc env env' -> f (env, t') t -> f (env', t') t
sink1 Extend ArrayR OpenAcc aenv aenv'
env1 OpenAcc aenv' brrs
OpenAcc (aenv, Array sh e) brrs
acc0
  = Text -> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a. Text -> a -> a
Stats.ruleFired Text
"aletD/eliminate"
  (Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs)
-> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a b. (a -> b) -> a -> b
$ case Cunctation r aenv' (Array sh e) -> Cunctation D aenv' (Array sh e)
forall r aenv sh e.
HasCallStack =>
Cunctation r aenv (Array sh e) -> Cunctation D aenv (Array sh e)
delaying Cunctation r aenv' arrs
Cunctation r aenv' (Array sh e)
cc1 of
      Step{}  -> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> OpenAcc (aenv', Array sh e) brrs
-> Embed OpenAcc aenv brrs
forall r aenv aenv' sh e brrs.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> OpenAcc (aenv', Array sh e) brrs
-> Embed OpenAcc aenv brrs
eliminate Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
Cunctation r aenv' (Array sh e)
cc1 OpenAcc (aenv', Array sh e) brrs
acc0'
      Yield{} -> Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> OpenAcc (aenv', Array sh e) brrs
-> Embed OpenAcc aenv brrs
forall r aenv aenv' sh e brrs.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> OpenAcc (aenv', Array sh e) brrs
-> Embed OpenAcc aenv brrs
eliminate Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
Cunctation r aenv' (Array sh e)
cc1 OpenAcc (aenv', Array sh e) brrs
acc0'

  where
    acc0 :: OpenAcc aenv' brrs
    acc0 :: OpenAcc aenv' brrs
acc0 = Embed OpenAcc aenv' brrs -> OpenAcc aenv' brrs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Extend ArrayR OpenAcc aenv' aenv'
-> Cunctation r aenv' brrs -> Embed OpenAcc aenv' brrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' brrs
cc0)

    kmap :: forall aenv a b. (PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b)
         -> OpenAcc aenv a
         -> OpenAcc aenv b
    kmap :: forall aenv a b.
(PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b)
-> OpenAcc aenv a -> OpenAcc aenv b
kmap PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b
f (OpenAcc PreOpenAcc OpenAcc aenv a
pacc) = PreOpenAcc OpenAcc aenv b -> OpenAcc aenv b
forall aenv t. PreOpenAcc OpenAcc aenv t -> OpenAcc aenv t
OpenAcc (PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b
f PreOpenAcc OpenAcc aenv a
pacc)

    -- The second part of let-elimination. Splitting into two steps exposes the
    -- extra type variables, and ensures we don't do extra work manipulating the
    -- body when not necessary (which can lead to a complexity blowup).
    --
    eliminate
        :: forall r aenv aenv' sh e brrs. HasCallStack
        => Extend ArrayR OpenAcc aenv aenv'
        -> Cunctation r aenv' (Array sh e)
        -> OpenAcc (aenv', Array sh e) brrs
        -> Embed OpenAcc aenv brrs
    eliminate :: forall r aenv aenv' sh e brrs.
HasCallStack =>
Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' (Array sh e)
-> OpenAcc (aenv', Array sh e) brrs
-> Embed OpenAcc aenv brrs
eliminate Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' (Array sh e)
cc1 OpenAcc (aenv', Array sh e) brrs
body
      | Done ArrayVars aenv' (Array sh e)
v1                  <- Cunctation r aenv' (Array sh e)
cc1
      , TupRsingle v1' :: Var ArrayR aenv' (Array sh e)
v1'@(Var ArrayR (Array sh e)
r Idx aenv' (Array sh e)
_) <- ArrayVars aenv' (Array sh e)
v1  = HasCallStack =>
ArrayR (Array sh e)
-> Exp aenv' sh -> Fun aenv' (sh -> e) -> Embed OpenAcc aenv brrs
ArrayR (Array sh e)
-> Exp aenv' sh -> Fun aenv' (sh -> e) -> Embed OpenAcc aenv brrs
elim ArrayR (Array sh e)
r (Var ArrayR aenv' (Array sh e) -> Exp aenv' sh
forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape Var ArrayR aenv' (Array sh e)
v1') (Var ArrayR aenv' (Array sh e) -> Fun aenv' (sh -> e)
forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray Var ArrayR aenv' (Array sh e)
v1')
      | Step ArrayR (Array sh' b)
r Exp aenv' sh'
sh1 Fun aenv' (sh' -> sh)
p1 Fun aenv' (a -> b)
f1 ArrayVar aenv' (Array sh a)
v1      <- Cunctation r aenv' (Array sh e)
cc1 = HasCallStack =>
ArrayR (Array sh e)
-> Exp aenv' sh -> Fun aenv' (sh -> e) -> Embed OpenAcc aenv brrs
ArrayR (Array sh e)
-> Exp aenv' sh -> Fun aenv' (sh -> e) -> Embed OpenAcc aenv brrs
elim ArrayR (Array sh e)
ArrayR (Array sh' b)
r Exp aenv' sh
Exp aenv' sh'
sh1 (OpenFun () aenv' (a -> e)
Fun aenv' (a -> b)
f1 OpenFun () aenv' (a -> e)
-> OpenFun () aenv' (sh -> a) -> Fun aenv' (sh -> e)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ArrayVar aenv' (Array sh a) -> Fun aenv' (sh -> a)
forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray ArrayVar aenv' (Array sh a)
v1 Fun aenv' (sh -> a)
-> OpenFun () aenv' (sh -> sh) -> OpenFun () aenv' (sh -> a)
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` OpenFun () aenv' (sh -> sh)
Fun aenv' (sh' -> sh)
p1)
      | Yield ArrayR (Array sh e)
r Exp aenv' sh
sh1 Fun aenv' (sh -> e)
f1           <- Cunctation r aenv' (Array sh e)
cc1 = HasCallStack =>
ArrayR (Array sh e)
-> Exp aenv' sh -> Fun aenv' (sh -> e) -> Embed OpenAcc aenv brrs
ArrayR (Array sh e)
-> Exp aenv' sh -> Fun aenv' (sh -> e) -> Embed OpenAcc aenv brrs
elim ArrayR (Array sh e)
ArrayR (Array sh e)
r Exp aenv' sh
Exp aenv' sh
sh1 Fun aenv' (sh -> e)
Fun aenv' (sh -> e)
f1
      where
        bnd :: PreOpenAcc OpenAcc aenv' (Array sh e)
        bnd :: PreOpenAcc OpenAcc aenv' (Array sh e)
bnd = Cunctation r aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sh e)
forall r aenv arrs.
HasCallStack =>
Cunctation r aenv arrs -> PreOpenAcc OpenAcc aenv arrs
compute Cunctation r aenv' (Array sh e)
cc1

        elim :: HasCallStack
             => ArrayR (Array sh e)
             -> Exp aenv' sh
             -> Fun aenv' (sh -> e)
             -> Embed OpenAcc aenv brrs
        elim :: HasCallStack =>
ArrayR (Array sh e)
-> Exp aenv' sh -> Fun aenv' (sh -> e) -> Embed OpenAcc aenv brrs
elim ArrayR (Array sh e)
r Exp aenv' sh
sh1 Fun aenv' (sh -> e)
f1
          | OpenExp () (aenv', Array sh e) sh
sh1'              <- (aenv' :> (aenv', Array sh e))
-> forall t. OpenExp () aenv' t -> OpenExp () (aenv', Array sh e) t
forall env env'.
(env :> env') -> forall t. OpenExp () env t -> OpenExp () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken ((aenv' :> aenv') -> aenv' :> (aenv', Array sh e)
forall env env' t. (env :> env') -> env :> (env', t)
weakenSucc' aenv' :> aenv'
forall env. env :> env
weakenId) Exp aenv' sh
sh1
          , OpenFun () (aenv', Array sh e) (sh -> e)
f1'               <- (aenv' :> (aenv', Array sh e))
-> forall t. OpenFun () aenv' t -> OpenFun () (aenv', Array sh e) t
forall env env'.
(env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken ((aenv' :> aenv') -> aenv' :> (aenv', Array sh e)
forall env env' t. (env :> env') -> env :> (env', t)
weakenSucc' aenv' :> aenv'
forall env. env :> env
weakenId) Fun aenv' (sh -> e)
f1
          , Embed Extend ArrayR OpenAcc aenv' aenv'
env0' Cunctation r aenv' brrs
cc0'  <- OpenAcc aenv' brrs -> Embed OpenAcc aenv' brrs
EmbedAcc OpenAcc
embedAcc (OpenAcc aenv' brrs -> Embed OpenAcc aenv' brrs)
-> OpenAcc aenv' brrs -> Embed OpenAcc aenv' brrs
forall a b. (a -> b) -> a -> b
$ (forall sh e.
 ArrayVar (aenv', Array sh e) (Array sh e)
 -> PreOpenAcc (AccClo OpenAcc) aenv' (Array sh e))
-> OpenAcc (aenv', Array sh e) brrs -> OpenAcc aenv' brrs
forall (f :: * -> * -> *) (fa :: (* -> * -> *) -> * -> * -> *) aenv
       aenv' a.
(Rebuildable f, SyntacticAcc fa) =>
(forall sh e.
 ArrayVar aenv (Array sh e) -> fa (AccClo f) aenv' (Array sh e))
-> f aenv a -> f aenv' a
forall (fa :: (* -> * -> *) -> * -> * -> *) aenv aenv' a.
SyntacticAcc fa =>
(forall sh e.
 ArrayVar aenv (Array sh e)
 -> fa (AccClo OpenAcc) aenv' (Array sh e))
-> OpenAcc aenv a -> OpenAcc aenv' a
rebuildA (PreOpenAcc OpenAcc aenv' (Array sh e)
-> Var ArrayR (aenv', Array sh e) (Array sh e)
-> PreOpenAcc OpenAcc aenv' (Array sh e)
forall (acc :: * -> * -> *) aenv t sh2 e2.
PreOpenAcc acc aenv t
-> ArrayVar (aenv, t) (Array sh2 e2)
-> PreOpenAcc acc aenv (Array sh2 e2)
subAtop PreOpenAcc OpenAcc aenv' (Array sh e)
bnd) (OpenAcc (aenv', Array sh e) brrs -> OpenAcc aenv' brrs)
-> OpenAcc (aenv', Array sh e) brrs -> OpenAcc aenv' brrs
forall a b. (a -> b) -> a -> b
$ (PreOpenAcc OpenAcc (aenv', Array sh e) brrs
 -> PreOpenAcc OpenAcc (aenv', Array sh e) brrs)
-> OpenAcc (aenv', Array sh e) brrs
-> OpenAcc (aenv', Array sh e) brrs
forall aenv a b.
(PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b)
-> OpenAcc aenv a -> OpenAcc aenv b
kmap (OpenExp () (aenv', Array sh e) sh
-> OpenFun () (aenv', Array sh e) (sh -> e)
-> ArrayVar (aenv', Array sh e) (Array sh e)
-> PreOpenAcc OpenAcc (aenv', Array sh e) brrs
-> PreOpenAcc OpenAcc (aenv', Array sh e) brrs
forall aenv sh e a.
HasCallStack =>
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
replaceA OpenExp () (aenv', Array sh e) sh
sh1' OpenFun () (aenv', Array sh e) (sh -> e)
f1' (ArrayVar (aenv', Array sh e) (Array sh e)
 -> PreOpenAcc OpenAcc (aenv', Array sh e) brrs
 -> PreOpenAcc OpenAcc (aenv', Array sh e) brrs)
-> ArrayVar (aenv', Array sh e) (Array sh e)
-> PreOpenAcc OpenAcc (aenv', Array sh e) brrs
-> PreOpenAcc OpenAcc (aenv', Array sh e) brrs
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e)
-> Idx (aenv', Array sh e) (Array sh e)
-> ArrayVar (aenv', Array sh e) (Array sh e)
forall (s :: * -> *) env t. s t -> Idx env t -> Var s env t
Var ArrayR (Array sh e)
r Idx (aenv', Array sh e) (Array sh e)
forall envt t env. (envt ~ (env, t)) => Idx envt t
ZeroIdx) OpenAcc (aenv', Array sh e) brrs
body
          = Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' brrs -> Embed OpenAcc aenv brrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv'
env1 Extend ArrayR OpenAcc aenv aenv'
-> Extend ArrayR OpenAcc aenv' aenv'
-> Extend ArrayR OpenAcc aenv aenv'
forall (s :: * -> *) (acc :: * -> * -> *) env env' env''.
Extend s acc env env'
-> Extend s acc env' env'' -> Extend s acc env env''
`append` Extend ArrayR OpenAcc aenv' aenv'
env0') Cunctation r aenv' brrs
cc0'

    -- As part of let-elimination, we need to replace uses of array variables in
    -- scalar expressions with an equivalent expression that generates the
    -- result directly
    --
    -- TODO: when we inline bindings we ought to let bind at the first
    --       occurrence and use a variable at all subsequent locations. At the
    --       moment we are just hoping CSE in the simplifier phase does good
    --       things, but that is limited in what it looks for.
    --
    replaceE :: forall env aenv sh e t. HasCallStack
             => OpenExp env aenv sh
             -> OpenFun env aenv (sh -> e)
             -> ArrayVar aenv (Array sh e)
             -> OpenExp env aenv t
             -> OpenExp env aenv t
    replaceE :: forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env aenv t
-> OpenExp env aenv t
replaceE OpenExp env aenv sh
sh' OpenFun env aenv (sh -> e)
f' avar :: ArrayVar aenv (Array sh e)
avar@(Var (ArrayR ShapeR sh
shR TypeR e
_) Idx aenv (Array sh e)
_) OpenExp env aenv t
exp =
      case OpenExp env aenv t
exp of
        Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
x OpenExp env' aenv t
y                     -> let k :: env :> env'
k = ELeftHandSide bnd_t env env' -> env :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ELeftHandSide bnd_t env env'
lhs
                                           in  ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide bnd_t env env'
lhs (OpenExp env aenv bnd_t -> OpenExp env aenv bnd_t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv bnd_t
x) (OpenExp env' aenv sh
-> OpenFun env' aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env' aenv t
-> OpenExp env' aenv t
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env aenv t
-> OpenExp env aenv t
replaceE ((env :> env') -> OpenExp env aenv sh -> OpenExp env' aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env :> env'
k OpenExp env aenv sh
sh') ((env :> env')
-> OpenFun env aenv (sh -> e) -> OpenFun env' aenv (sh -> e)
forall env env' aenv t.
(env :> env') -> OpenFun env aenv t -> OpenFun env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env :> env'
k OpenFun env aenv (sh -> e)
f') ArrayVar aenv (Array sh e)
avar OpenExp env' aenv t
y)
        Evar ExpVar env t
var                        -> ExpVar env t -> OpenExp env aenv t
forall env t aenv. ExpVar env t -> OpenExp env aenv t
Evar ExpVar env t
var
        Foreign TypeR t
tR asm (x -> t)
ff Fun () (x -> t)
f OpenExp env aenv x
e               -> TypeR t
-> asm (x -> t)
-> Fun () (x -> t)
-> OpenExp env aenv x
-> OpenExp env aenv t
forall (asm :: * -> *) t x env aenv.
Foreign asm =>
TypeR t
-> asm (x -> t)
-> Fun () (x -> t)
-> OpenExp env aenv x
-> OpenExp env aenv t
Foreign TypeR t
tR asm (x -> t)
ff Fun () (x -> t)
f (OpenExp env aenv x -> OpenExp env aenv x
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv x
e)
        Const ScalarType t
tR t
c                      -> ScalarType t -> t -> OpenExp env aenv t
forall t env aenv. ScalarType t -> t -> OpenExp env aenv t
Const ScalarType t
tR t
c
        Undef ScalarType t
tR                        -> ScalarType t -> OpenExp env aenv t
forall t env aenv. ScalarType t -> OpenExp env aenv t
Undef ScalarType t
tR
        OpenExp env aenv t
Nil                             -> OpenExp env aenv t
OpenExp env aenv ()
forall env aenv. OpenExp env aenv ()
Nil
        Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2                      -> OpenExp env aenv t1
-> OpenExp env aenv t2 -> OpenExp env aenv (t1, t2)
forall env aenv t1 t2.
OpenExp env aenv t1
-> OpenExp env aenv t2 -> OpenExp env aenv (t1, t2)
Pair (OpenExp env aenv t1 -> OpenExp env aenv t1
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv t1
e1) (OpenExp env aenv t2 -> OpenExp env aenv t2
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv t2
e2)
        VecPack VecR n s tup
vR OpenExp env aenv tup
e                    -> VecR n s tup -> OpenExp env aenv tup -> OpenExp env aenv (Vec n s)
forall (n :: Nat) s tup env aenv.
KnownNat n =>
VecR n s tup -> OpenExp env aenv tup -> OpenExp env aenv (Vec n s)
VecPack VecR n s tup
vR (OpenExp env aenv tup -> OpenExp env aenv tup
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv tup
e)
        VecUnpack VecR n s t
vR OpenExp env aenv (Vec n s)
e                  -> VecR n s t -> OpenExp env aenv (Vec n s) -> OpenExp env aenv t
forall (n :: Nat) s t env aenv.
KnownNat n =>
VecR n s t -> OpenExp env aenv (Vec n s) -> OpenExp env aenv t
VecUnpack VecR n s t
vR (OpenExp env aenv (Vec n s) -> OpenExp env aenv (Vec n s)
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv (Vec n s)
e)
        IndexSlice SliceIndex slix t co sh
x OpenExp env aenv slix
ix OpenExp env aenv sh
sh              -> SliceIndex slix t co sh
-> OpenExp env aenv slix
-> OpenExp env aenv sh
-> OpenExp env aenv t
forall slix t co sh env aenv.
SliceIndex slix t co sh
-> OpenExp env aenv slix
-> OpenExp env aenv sh
-> OpenExp env aenv t
IndexSlice SliceIndex slix t co sh
x (OpenExp env aenv slix -> OpenExp env aenv slix
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv slix
ix) (OpenExp env aenv sh -> OpenExp env aenv sh
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv sh
sh)
        IndexFull SliceIndex slix sl co t
x OpenExp env aenv slix
ix OpenExp env aenv sl
sl               -> SliceIndex slix sl co t
-> OpenExp env aenv slix
-> OpenExp env aenv sl
-> OpenExp env aenv t
forall slix sl co t env aenv.
SliceIndex slix sl co t
-> OpenExp env aenv slix
-> OpenExp env aenv sl
-> OpenExp env aenv t
IndexFull SliceIndex slix sl co t
x (OpenExp env aenv slix -> OpenExp env aenv slix
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv slix
ix) (OpenExp env aenv sl -> OpenExp env aenv sl
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv sl
sl)
        ToIndex ShapeR sh
shR' OpenExp env aenv sh
sh OpenExp env aenv sh
ix              -> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv Int
forall sh env aenv.
ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv Int
ToIndex ShapeR sh
shR' (OpenExp env aenv sh -> OpenExp env aenv sh
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv sh
sh) (OpenExp env aenv sh -> OpenExp env aenv sh
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv sh
ix)
        FromIndex ShapeR t
shR' OpenExp env aenv t
sh OpenExp env aenv Int
i             -> ShapeR t
-> OpenExp env aenv t -> OpenExp env aenv Int -> OpenExp env aenv t
forall t env aenv.
ShapeR t
-> OpenExp env aenv t -> OpenExp env aenv Int -> OpenExp env aenv t
FromIndex ShapeR t
shR' (OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv t
sh) (OpenExp env aenv Int -> OpenExp env aenv Int
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv Int
i)
        Case OpenExp env aenv PrimBool
e [(PrimBool, OpenExp env aenv t)]
rhs Maybe (OpenExp env aenv t)
def                  -> OpenExp env aenv PrimBool
-> [(PrimBool, OpenExp env aenv t)]
-> Maybe (OpenExp env aenv t)
-> OpenExp env aenv t
forall env aenv t.
OpenExp env aenv PrimBool
-> [(PrimBool, OpenExp env aenv t)]
-> Maybe (OpenExp env aenv t)
-> OpenExp env aenv t
Case (OpenExp env aenv PrimBool -> OpenExp env aenv PrimBool
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv PrimBool
e) (ASetter
  [(PrimBool, OpenExp env aenv t)]
  [(PrimBool, OpenExp env aenv t)]
  (OpenExp env aenv t)
  (OpenExp env aenv t)
-> (OpenExp env aenv t -> OpenExp env aenv t)
-> [(PrimBool, OpenExp env aenv t)]
-> [(PrimBool, OpenExp env aenv t)]
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (ASetter
  [(PrimBool, OpenExp env aenv t)]
  [(PrimBool, OpenExp env aenv t)]
  (PrimBool, OpenExp env aenv t)
  (PrimBool, OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => ASetter (f a) (f b) a b
mapped ASetter
  [(PrimBool, OpenExp env aenv t)]
  [(PrimBool, OpenExp env aenv t)]
  (PrimBool, OpenExp env aenv t)
  (PrimBool, OpenExp env aenv t)
-> ((OpenExp env aenv t -> Identity (OpenExp env aenv t))
    -> (PrimBool, OpenExp env aenv t)
    -> Identity (PrimBool, OpenExp env aenv t))
-> ASetter
     [(PrimBool, OpenExp env aenv t)]
     [(PrimBool, OpenExp env aenv t)]
     (OpenExp env aenv t)
     (OpenExp env aenv t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (OpenExp env aenv t -> Identity (OpenExp env aenv t))
-> (PrimBool, OpenExp env aenv t)
-> Identity (PrimBool, OpenExp env aenv t)
forall s t a b. Field2 s t a b => Lens s t a b
Lens
  (PrimBool, OpenExp env aenv t)
  (PrimBool, OpenExp env aenv t)
  (OpenExp env aenv t)
  (OpenExp env aenv t)
_2) OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE [(PrimBool, OpenExp env aenv t)]
rhs) ((OpenExp env aenv t -> OpenExp env aenv t)
-> Maybe (OpenExp env aenv t) -> Maybe (OpenExp env aenv t)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE Maybe (OpenExp env aenv t)
def)
        Cond OpenExp env aenv PrimBool
p OpenExp env aenv t
t OpenExp env aenv t
e                      -> OpenExp env aenv PrimBool
-> OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t
forall env aenv t.
OpenExp env aenv PrimBool
-> OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t
Cond (OpenExp env aenv PrimBool -> OpenExp env aenv PrimBool
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv PrimBool
p) (OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv t
t) (OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv t
e)
        PrimConst PrimConst t
c                     -> PrimConst t -> OpenExp env aenv t
forall t env aenv. PrimConst t -> OpenExp env aenv t
PrimConst PrimConst t
c
        PrimApp PrimFun (a -> t)
g OpenExp env aenv a
x                     -> PrimFun (a -> t) -> OpenExp env aenv a -> OpenExp env aenv t
forall a t env aenv.
PrimFun (a -> t) -> OpenExp env aenv a -> OpenExp env aenv t
PrimApp PrimFun (a -> t)
g (OpenExp env aenv a -> OpenExp env aenv a
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv a
x)
        ShapeSize ShapeR dim
shR' OpenExp env aenv dim
sh               -> ShapeR dim -> OpenExp env aenv dim -> OpenExp env aenv Int
forall dim env aenv.
ShapeR dim -> OpenExp env aenv dim -> OpenExp env aenv Int
ShapeSize ShapeR dim
shR' (OpenExp env aenv dim -> OpenExp env aenv dim
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv dim
sh)
        While OpenFun env aenv (t -> PrimBool)
p OpenFun env aenv (t -> t)
f OpenExp env aenv t
x                     -> OpenFun env aenv (t -> PrimBool)
-> OpenFun env aenv (t -> t)
-> OpenExp env aenv t
-> OpenExp env aenv t
forall env aenv t.
OpenFun env aenv (t -> PrimBool)
-> OpenFun env aenv (t -> t)
-> OpenExp env aenv t
-> OpenExp env aenv t
While (OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env aenv (t -> PrimBool)
-> OpenFun env aenv (t -> PrimBool)
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env aenv t
-> OpenFun env aenv t
replaceF OpenExp env aenv sh
sh' OpenFun env aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar OpenFun env aenv (t -> PrimBool)
p) (OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env aenv (t -> t)
-> OpenFun env aenv (t -> t)
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env aenv t
-> OpenFun env aenv t
replaceF OpenExp env aenv sh
sh' OpenFun env aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar OpenFun env aenv (t -> t)
f) (OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv t
x)
        Coerce ScalarType a
t1 ScalarType t
t2 OpenExp env aenv a
e                  -> ScalarType a
-> ScalarType t -> OpenExp env aenv a -> OpenExp env aenv t
forall a t env aenv.
BitSizeEq a t =>
ScalarType a
-> ScalarType t -> OpenExp env aenv a -> OpenExp env aenv t
Coerce ScalarType a
t1 ScalarType t
t2 (OpenExp env aenv a -> OpenExp env aenv a
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv a
e)

        Shape ArrayVar aenv (Array t e)
a
          | Just Array t e :~: Array sh e
Refl <- ArrayVar aenv (Array t e)
-> ArrayVar aenv (Array sh e) -> Maybe (Array t e :~: Array sh e)
forall (s :: * -> *) env t1 t2.
Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2)
matchVar ArrayVar aenv (Array t e)
a ArrayVar aenv (Array sh e)
avar -> Text -> OpenExp env aenv t -> OpenExp env aenv t
forall a. Text -> a -> a
Stats.substitution Text
"replaceE/shape" OpenExp env aenv sh
OpenExp env aenv t
sh'
          | Bool
otherwise                    -> OpenExp env aenv t
exp

        Index ArrayVar aenv (Array dim t)
a OpenExp env aenv dim
sh
          | Just Array dim t :~: Array sh e
Refl        <- ArrayVar aenv (Array dim t)
-> ArrayVar aenv (Array sh e) -> Maybe (Array dim t :~: Array sh e)
forall (s :: * -> *) env t1 t2.
Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2)
matchVar ArrayVar aenv (Array dim t)
a ArrayVar aenv (Array sh e)
avar
          , Lam ELeftHandSide a env env'
lhs (Body OpenExp env' aenv t1
b) <- OpenFun env aenv (sh -> e)
f'      -> Text -> OpenExp env aenv t -> OpenExp env aenv t
forall a. Text -> a -> a
Stats.substitution Text
"replaceE/!" (OpenExp env aenv t -> OpenExp env aenv t)
-> (OpenExp env aenv t -> OpenExp env aenv t)
-> OpenExp env aenv t
-> OpenExp env aenv t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE (OpenExp env aenv t -> OpenExp env aenv t)
-> OpenExp env aenv t -> OpenExp env aenv t
forall a b. (a -> b) -> a -> b
$ ELeftHandSide a env env'
-> OpenExp env aenv a -> OpenExp env' aenv t -> OpenExp env aenv t
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide a env env'
lhs OpenExp env aenv dim
OpenExp env aenv a
sh OpenExp env' aenv t
OpenExp env' aenv t1
b
          | Bool
otherwise                   -> ArrayVar aenv (Array dim t)
-> OpenExp env aenv dim -> OpenExp env aenv t
forall aenv dim t env.
ArrayVar aenv (Array dim t)
-> OpenExp env aenv dim -> OpenExp env aenv t
Index ArrayVar aenv (Array dim t)
a (OpenExp env aenv dim -> OpenExp env aenv dim
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv dim
sh)

        LinearIndex ArrayVar aenv (Array dim t)
a OpenExp env aenv Int
i
          | Just Array dim t :~: Array sh e
Refl        <- ArrayVar aenv (Array dim t)
-> ArrayVar aenv (Array sh e) -> Maybe (Array dim t :~: Array sh e)
forall (s :: * -> *) env t1 t2.
Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2)
matchVar ArrayVar aenv (Array dim t)
a ArrayVar aenv (Array sh e)
avar
          , Lam ELeftHandSide a env env'
lhs (Body OpenExp env' aenv t1
b) <- OpenFun env aenv (sh -> e)
f'
                                        -> Text -> OpenExp env aenv t -> OpenExp env aenv t
forall a. Text -> a -> a
Stats.substitution Text
"replaceE/!!" (OpenExp env aenv t -> OpenExp env aenv t)
-> (OpenExp env aenv t -> OpenExp env aenv t)
-> OpenExp env aenv t
-> OpenExp env aenv t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenExp env aenv t -> OpenExp env aenv t
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE
                                         (OpenExp env aenv t -> OpenExp env aenv t)
-> OpenExp env aenv t -> OpenExp env aenv t
forall a b. (a -> b) -> a -> b
$ ELeftHandSide a env env'
-> OpenExp env aenv a -> OpenExp env' aenv t -> OpenExp env aenv t
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide a env env'
lhs
                                               (ELeftHandSide Int env (env, Int)
-> OpenExp env aenv Int
-> OpenExp (env, Int) aenv a
-> OpenExp env aenv a
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let (ScalarType Int -> ELeftHandSide Int env (env, Int)
forall (s :: * -> *) v env. s v -> LeftHandSide s v env (env, v)
LeftHandSideSingle ScalarType Int
scalarTypeInt) OpenExp env aenv Int
i
                                                    (OpenExp (env, Int) aenv a -> OpenExp env aenv a)
-> OpenExp (env, Int) aenv a -> OpenExp env aenv a
forall a b. (a -> b) -> a -> b
$ ShapeR a
-> OpenExp (env, Int) aenv a
-> OpenExp (env, Int) aenv Int
-> OpenExp (env, Int) aenv a
forall t env aenv.
ShapeR t
-> OpenExp env aenv t -> OpenExp env aenv Int -> OpenExp env aenv t
FromIndex ShapeR sh
ShapeR a
shR ((env :> (env, Int))
-> OpenExp env aenv a -> OpenExp (env, Int) aenv a
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE ((env :> env) -> env :> (env, Int)
forall env env' t. (env :> env') -> env :> (env', t)
weakenSucc' env :> env
forall env. env :> env
weakenId) OpenExp env aenv sh
OpenExp env aenv a
sh')
                                                    (OpenExp (env, Int) aenv Int -> OpenExp (env, Int) aenv a)
-> OpenExp (env, Int) aenv Int -> OpenExp (env, Int) aenv a
forall a b. (a -> b) -> a -> b
$ ExpVar (env, Int) Int -> OpenExp (env, Int) aenv Int
forall env t aenv. ExpVar env t -> OpenExp env aenv t
Evar
                                                    (ExpVar (env, Int) Int -> OpenExp (env, Int) aenv Int)
-> ExpVar (env, Int) Int -> OpenExp (env, Int) aenv Int
forall a b. (a -> b) -> a -> b
$ ScalarType Int -> Idx (env, Int) Int -> ExpVar (env, Int) Int
forall (s :: * -> *) env t. s t -> Idx env t -> Var s env t
Var ScalarType Int
scalarTypeInt Idx (env, Int) Int
forall envt t env. (envt ~ (env, t)) => Idx envt t
ZeroIdx)
                                               OpenExp env' aenv t
OpenExp env' aenv t1
b
          | Bool
otherwise                   -> ArrayVar aenv (Array dim t)
-> OpenExp env aenv Int -> OpenExp env aenv t
forall aenv dim t env.
ArrayVar aenv (Array dim t)
-> OpenExp env aenv Int -> OpenExp env aenv t
LinearIndex ArrayVar aenv (Array dim t)
a (OpenExp env aenv Int -> OpenExp env aenv Int
forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE OpenExp env aenv Int
i)

      where
        cvtE :: OpenExp env aenv s -> OpenExp env aenv s
        cvtE :: forall s. OpenExp env aenv s -> OpenExp env aenv s
cvtE = OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env aenv s
-> OpenExp env aenv s
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env aenv t
-> OpenExp env aenv t
replaceE OpenExp env aenv sh
sh' OpenFun env aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar

    replaceF :: forall env aenv sh e t. HasCallStack
             => OpenExp env aenv sh
             -> OpenFun env aenv (sh -> e)
             -> ArrayVar aenv (Array sh e)
             -> OpenFun env aenv t
             -> OpenFun env aenv t
    replaceF :: forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env aenv t
-> OpenFun env aenv t
replaceF OpenExp env aenv sh
sh' OpenFun env aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar OpenFun env aenv t
fun =
      case OpenFun env aenv t
fun of
        Body OpenExp env aenv t
e          -> OpenExp env aenv t -> OpenFun env aenv t
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env aenv t
-> OpenExp env aenv t
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env aenv t
-> OpenExp env aenv t
replaceE OpenExp env aenv sh
sh' OpenFun env aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar OpenExp env aenv t
e)
        Lam ELeftHandSide a env env'
lhs OpenFun env' aenv t1
f       -> let k :: env :> env'
k = ELeftHandSide a env env' -> env :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ELeftHandSide a env env'
lhs
                           in  ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam ELeftHandSide a env env'
lhs (OpenExp env' aenv sh
-> OpenFun env' aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env' aenv t1
-> OpenFun env' aenv t1
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env aenv t
-> OpenFun env aenv t
replaceF ((env :> env') -> OpenExp env aenv sh -> OpenExp env' aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env :> env'
k OpenExp env aenv sh
sh') ((env :> env')
-> OpenFun env aenv (sh -> e) -> OpenFun env' aenv (sh -> e)
forall env env' aenv t.
(env :> env') -> OpenFun env aenv t -> OpenFun env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env :> env'
k OpenFun env aenv (sh -> e)
f') ArrayVar aenv (Array sh e)
avar OpenFun env' aenv t1
f)

    replaceA :: forall aenv sh e a. HasCallStack
             => Exp aenv sh
             -> Fun aenv (sh -> e)
             -> ArrayVar aenv (Array sh e)
             -> PreOpenAcc OpenAcc aenv a
             -> PreOpenAcc OpenAcc aenv a
    replaceA :: forall aenv sh e a.
HasCallStack =>
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
replaceA Exp aenv sh
sh' Fun aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar PreOpenAcc OpenAcc aenv a
pacc =
      case PreOpenAcc OpenAcc aenv a
pacc of
        Avar ArrayVar aenv (Array sh e)
v
          | Just Array sh e :~: Array sh e
Refl <- ArrayVar aenv (Array sh e)
-> ArrayVar aenv (Array sh e) -> Maybe (Array sh e :~: Array sh e)
forall (s :: * -> *) env t1 t2.
Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2)
matchVar ArrayVar aenv (Array sh e)
v ArrayVar aenv (Array sh e)
avar -> ArrayVar aenv (Array sh e) -> PreOpenAcc OpenAcc aenv (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar ArrayVar aenv (Array sh e)
avar
          | Bool
otherwise                    -> ArrayVar aenv (Array sh e) -> PreOpenAcc OpenAcc aenv (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar ArrayVar aenv (Array sh e)
v

        Alet ALeftHandSide bndArrs aenv aenv'
lhs OpenAcc aenv bndArrs
bnd (OpenAcc aenv' a
body :: OpenAcc aenv1 a) ->
          let w :: aenv :> aenv1
              w :: aenv :> aenv'
w    = ALeftHandSide bndArrs aenv aenv' -> aenv :> aenv'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ALeftHandSide bndArrs aenv aenv'
lhs
              sh'' :: OpenExp () aenv' sh
sh'' = (aenv :> aenv')
-> forall t. OpenExp () aenv t -> OpenExp () aenv' t
forall env env'.
(env :> env') -> forall t. OpenExp () env t -> OpenExp () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken aenv :> aenv'
w Exp aenv sh
sh'
              f'' :: OpenFun () aenv' (sh -> e)
f''  = (aenv :> aenv')
-> forall t. OpenFun () aenv t -> OpenFun () aenv' t
forall env env'.
(env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken aenv :> aenv'
w Fun aenv (sh -> e)
f'
          in
          ALeftHandSide bndArrs aenv aenv'
-> OpenAcc aenv bndArrs
-> OpenAcc aenv' a
-> PreOpenAcc OpenAcc aenv a
forall bndArrs aenv aenv' (acc :: * -> * -> *) a.
ALeftHandSide bndArrs aenv aenv'
-> acc aenv bndArrs -> acc aenv' a -> PreOpenAcc acc aenv a
Alet ALeftHandSide bndArrs aenv aenv'
lhs (OpenAcc aenv bndArrs -> OpenAcc aenv bndArrs
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv bndArrs
bnd) ((PreOpenAcc OpenAcc aenv' a -> PreOpenAcc OpenAcc aenv' a)
-> OpenAcc aenv' a -> OpenAcc aenv' a
forall aenv a b.
(PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b)
-> OpenAcc aenv a -> OpenAcc aenv b
kmap (OpenExp () aenv' sh
-> OpenFun () aenv' (sh -> e)
-> ArrayVar aenv' (Array sh e)
-> PreOpenAcc OpenAcc aenv' a
-> PreOpenAcc OpenAcc aenv' a
forall aenv sh e a.
HasCallStack =>
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
replaceA OpenExp () aenv' sh
sh'' OpenFun () aenv' (sh -> e)
f'' ((aenv :> aenv')
-> forall t. Var ArrayR aenv t -> Var ArrayR aenv' t
forall env env'.
(env :> env') -> forall t. Var ArrayR env t -> Var ArrayR env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken aenv :> aenv'
w ArrayVar aenv (Array sh e)
avar)) OpenAcc aenv' a
body)

        Use ArrayR (Array sh e)
repr Array sh e
arrs           -> ArrayR (Array sh e)
-> Array sh e -> PreOpenAcc OpenAcc aenv (Array sh e)
forall sh e (acc :: * -> * -> *) aenv.
ArrayR (Array sh e)
-> Array sh e -> PreOpenAcc acc aenv (Array sh e)
Use ArrayR (Array sh e)
repr Array sh e
arrs
        Unit TypeR e
tR Exp aenv e
e               -> TypeR e -> Exp aenv e -> PreOpenAcc OpenAcc aenv (Array () e)
forall e aenv (acc :: * -> * -> *).
TypeR e -> Exp aenv e -> PreOpenAcc acc aenv (Array () e)
Unit TypeR e
tR (Exp aenv e -> Exp aenv e
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv e
e)
        Acond Exp aenv PrimBool
p OpenAcc aenv a
at OpenAcc aenv a
ae           -> Exp aenv PrimBool
-> OpenAcc aenv a -> OpenAcc aenv a -> PreOpenAcc OpenAcc aenv a
forall aenv (acc :: * -> * -> *) a.
Exp aenv PrimBool
-> acc aenv a -> acc aenv a -> PreOpenAcc acc aenv a
Acond (Exp aenv PrimBool -> Exp aenv PrimBool
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv PrimBool
p) (OpenAcc aenv a -> OpenAcc aenv a
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv a
at) (OpenAcc aenv a -> OpenAcc aenv a
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv a
ae)
        PreOpenAcc OpenAcc aenv a
Anil                    -> PreOpenAcc OpenAcc aenv a
PreOpenAcc OpenAcc aenv ()
forall (acc :: * -> * -> *) aenv. PreOpenAcc acc aenv ()
Anil
        Atrace Message arrs1
msg OpenAcc aenv arrs1
a OpenAcc aenv a
b          -> Message arrs1
-> OpenAcc aenv arrs1
-> OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
forall arrs1 (acc :: * -> * -> *) aenv a.
Message arrs1
-> acc aenv arrs1 -> acc aenv a -> PreOpenAcc acc aenv a
Atrace Message arrs1
msg (OpenAcc aenv arrs1 -> OpenAcc aenv arrs1
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv arrs1
a) (OpenAcc aenv a -> OpenAcc aenv a
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv a
b)
        Apair OpenAcc aenv as
a1 OpenAcc aenv bs
a2             -> OpenAcc aenv as
-> OpenAcc aenv bs -> PreOpenAcc OpenAcc aenv (as, bs)
forall (acc :: * -> * -> *) aenv as bs.
acc aenv as -> acc aenv bs -> PreOpenAcc acc aenv (as, bs)
Apair (OpenAcc aenv as -> OpenAcc aenv as
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv as
a1) (OpenAcc aenv bs -> OpenAcc aenv bs
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv bs
a2)
        Awhile PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
p PreOpenAfun OpenAcc aenv (a -> a)
f OpenAcc aenv a
a            -> PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
-> PreOpenAfun OpenAcc aenv (a -> a)
-> OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
forall (acc :: * -> * -> *) aenv a.
PreOpenAfun acc aenv (a -> Scalar PrimBool)
-> PreOpenAfun acc aenv (a -> a)
-> acc aenv a
-> PreOpenAcc acc aenv a
Awhile (PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
-> PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
forall s.
HasCallStack =>
PreOpenAfun OpenAcc aenv s -> PreOpenAfun OpenAcc aenv s
cvtAF PreOpenAfun OpenAcc aenv (a -> Scalar PrimBool)
p) (PreOpenAfun OpenAcc aenv (a -> a)
-> PreOpenAfun OpenAcc aenv (a -> a)
forall s.
HasCallStack =>
PreOpenAfun OpenAcc aenv s -> PreOpenAfun OpenAcc aenv s
cvtAF PreOpenAfun OpenAcc aenv (a -> a)
f) (OpenAcc aenv a -> OpenAcc aenv a
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv a
a)
        Apply ArraysR a
repr PreOpenAfun OpenAcc aenv (arrs1 -> a)
f OpenAcc aenv arrs1
a          -> ArraysR a
-> PreOpenAfun OpenAcc aenv (arrs1 -> a)
-> OpenAcc aenv arrs1
-> PreOpenAcc OpenAcc aenv a
forall a (acc :: * -> * -> *) aenv arrs1.
ArraysR a
-> PreOpenAfun acc aenv (arrs1 -> a)
-> acc aenv arrs1
-> PreOpenAcc acc aenv a
Apply ArraysR a
repr (PreOpenAfun OpenAcc aenv (arrs1 -> a)
-> PreOpenAfun OpenAcc aenv (arrs1 -> a)
forall s.
HasCallStack =>
PreOpenAfun OpenAcc aenv s -> PreOpenAfun OpenAcc aenv s
cvtAF PreOpenAfun OpenAcc aenv (arrs1 -> a)
f) (OpenAcc aenv arrs1 -> OpenAcc aenv arrs1
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv arrs1
a)
        Aforeign ArraysR a
repr asm (as -> a)
ff PreAfun OpenAcc (as -> a)
f OpenAcc aenv as
a    -> ArraysR a
-> asm (as -> a)
-> PreAfun OpenAcc (as -> a)
-> OpenAcc aenv as
-> PreOpenAcc OpenAcc aenv a
forall (asm :: * -> *) a as (acc :: * -> * -> *) aenv.
Foreign asm =>
ArraysR a
-> asm (as -> a)
-> PreAfun acc (as -> a)
-> acc aenv as
-> PreOpenAcc acc aenv a
Aforeign ArraysR a
repr asm (as -> a)
ff PreAfun OpenAcc (as -> a)
f (OpenAcc aenv as -> OpenAcc aenv as
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv as
a)       -- no sharing between f and a
        Generate ArrayR (Array sh e)
repr Exp aenv sh
sh Fun aenv (sh -> e)
f      -> ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> PreOpenAcc OpenAcc aenv (Array sh e)
forall sh e aenv (acc :: * -> * -> *).
ArrayR (Array sh e)
-> Exp aenv sh
-> Fun aenv (sh -> e)
-> PreOpenAcc acc aenv (Array sh e)
Generate ArrayR (Array sh e)
repr (Exp aenv sh -> Exp aenv sh
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv sh
sh) (Fun aenv (sh -> e) -> Fun aenv (sh -> e)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (sh -> e)
f)
        Map TypeR e'
tR Fun aenv (e -> e')
f OpenAcc aenv (Array sh e)
a              -> TypeR e'
-> Fun aenv (e -> e')
-> OpenAcc aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv (Array sh e')
forall e' aenv e (acc :: * -> * -> *) sh.
TypeR e'
-> Fun aenv (e -> e')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Map TypeR e'
tR (Fun aenv (e -> e') -> Fun aenv (e -> e')
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (e -> e')
f) (OpenAcc aenv (Array sh e) -> OpenAcc aenv (Array sh e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh e)
a)
        ZipWith TypeR e3
tR Fun aenv (e1 -> e2 -> e3)
f OpenAcc aenv (Array sh e1)
a OpenAcc aenv (Array sh e2)
b        -> TypeR e3
-> Fun aenv (e1 -> e2 -> e3)
-> OpenAcc aenv (Array sh e1)
-> OpenAcc aenv (Array sh e2)
-> PreOpenAcc OpenAcc aenv (Array sh e3)
forall e3 aenv e1 e2 (acc :: * -> * -> *) sh.
TypeR e3
-> Fun aenv (e1 -> e2 -> e3)
-> acc aenv (Array sh e1)
-> acc aenv (Array sh e2)
-> PreOpenAcc acc aenv (Array sh e3)
ZipWith TypeR e3
tR (Fun aenv (e1 -> e2 -> e3) -> Fun aenv (e1 -> e2 -> e3)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (e1 -> e2 -> e3)
f) (OpenAcc aenv (Array sh e1) -> OpenAcc aenv (Array sh e1)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh e1)
a) (OpenAcc aenv (Array sh e2) -> OpenAcc aenv (Array sh e2)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh e2)
b)
        Backpermute ShapeR sh'
shR Exp aenv sh'
sh Fun aenv (sh' -> sh)
p OpenAcc aenv (Array sh e)
a  -> ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> OpenAcc aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv (Array sh' e)
forall sh' aenv sh (acc :: * -> * -> *) e.
ShapeR sh'
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Backpermute ShapeR sh'
shR (Exp aenv sh' -> Exp aenv sh'
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv sh'
sh) (Fun aenv (sh' -> sh) -> Fun aenv (sh' -> sh)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (sh' -> sh)
p) (OpenAcc aenv (Array sh e) -> OpenAcc aenv (Array sh e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh e)
a)
        Transform ArrayR (Array sh' b)
repr Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a1 -> b)
f OpenAcc aenv (Array sh a1)
a -> ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> OpenAcc aenv (Array sh a1)
-> PreOpenAcc OpenAcc aenv (Array sh' b)
forall sh' b aenv sh a1 (acc :: * -> * -> *).
ArrayR (Array sh' b)
-> Exp aenv sh'
-> Fun aenv (sh' -> sh)
-> Fun aenv (a1 -> b)
-> acc aenv (Array sh a1)
-> PreOpenAcc acc aenv (Array sh' b)
Transform ArrayR (Array sh' b)
repr (Exp aenv sh' -> Exp aenv sh'
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv sh'
sh) (Fun aenv (sh' -> sh) -> Fun aenv (sh' -> sh)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (sh' -> sh)
p) (Fun aenv (a1 -> b) -> Fun aenv (a1 -> b)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (a1 -> b)
f) (OpenAcc aenv (Array sh a1) -> OpenAcc aenv (Array sh a1)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh a1)
a)
        Slice SliceIndex slix sl co sh
slix OpenAcc aenv (Array sh e)
a Exp aenv slix
sl         -> SliceIndex slix sl co sh
-> OpenAcc aenv (Array sh e)
-> Exp aenv slix
-> PreOpenAcc OpenAcc aenv (Array sl e)
forall slix sl co sh (acc :: * -> * -> *) aenv e.
SliceIndex slix sl co sh
-> acc aenv (Array sh e)
-> Exp aenv slix
-> PreOpenAcc acc aenv (Array sl e)
Slice SliceIndex slix sl co sh
slix (OpenAcc aenv (Array sh e) -> OpenAcc aenv (Array sh e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh e)
a) (Exp aenv slix -> Exp aenv slix
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv slix
sl)
        Replicate SliceIndex slix sl co sh
slix Exp aenv slix
sh OpenAcc aenv (Array sl e)
a     -> SliceIndex slix sl co sh
-> Exp aenv slix
-> OpenAcc aenv (Array sl e)
-> PreOpenAcc OpenAcc aenv (Array sh e)
forall slix sl co sh aenv (acc :: * -> * -> *) e.
SliceIndex slix sl co sh
-> Exp aenv slix
-> acc aenv (Array sl e)
-> PreOpenAcc acc aenv (Array sh e)
Replicate SliceIndex slix sl co sh
slix (Exp aenv slix -> Exp aenv slix
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv slix
sh) (OpenAcc aenv (Array sl e) -> OpenAcc aenv (Array sl e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sl e)
a)
        Reshape ShapeR sh
shR Exp aenv sh
sl OpenAcc aenv (Array sh' e)
a        -> ShapeR sh
-> Exp aenv sh
-> OpenAcc aenv (Array sh' e)
-> PreOpenAcc OpenAcc aenv (Array sh e)
forall sh aenv (acc :: * -> * -> *) sh' e.
ShapeR sh
-> Exp aenv sh
-> acc aenv (Array sh' e)
-> PreOpenAcc acc aenv (Array sh e)
Reshape ShapeR sh
shR (Exp aenv sh -> Exp aenv sh
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv sh
sl) (OpenAcc aenv (Array sh' e) -> OpenAcc aenv (Array sh' e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh' e)
a)
        Fold Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a              -> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> OpenAcc aenv (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv (Array sh e)
forall aenv e (acc :: * -> * -> *) sh.
Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array sh e)
Fold (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall s. Exp aenv s -> Exp aenv s
cvtE (Exp aenv e -> Exp aenv e)
-> Maybe (Exp aenv e) -> Maybe (Exp aenv e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Exp aenv e)
z) (OpenAcc aenv (Array (sh, Int) e)
-> OpenAcc aenv (Array (sh, Int) e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array (sh, Int) e)
a)
        FoldSeg IntegralType i
i Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a OpenAcc aenv (Segments i)
s       -> IntegralType i
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> OpenAcc aenv (Array (sh, Int) e)
-> OpenAcc aenv (Segments i)
-> PreOpenAcc OpenAcc aenv (Array (sh, Int) e)
forall i aenv e (acc :: * -> * -> *) sh.
IntegralType i
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> acc aenv (Segments i)
-> PreOpenAcc acc aenv (Array (sh, Int) e)
FoldSeg IntegralType i
i (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall s. Exp aenv s -> Exp aenv s
cvtE (Exp aenv e -> Exp aenv e)
-> Maybe (Exp aenv e) -> Maybe (Exp aenv e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Exp aenv e)
z) (OpenAcc aenv (Array (sh, Int) e)
-> OpenAcc aenv (Array (sh, Int) e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array (sh, Int) e)
a) (OpenAcc aenv (Segments i) -> OpenAcc aenv (Segments i)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Segments i)
s)
        Scan  Direction
d Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z OpenAcc aenv (Array (sh, Int) e)
a           -> Direction
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> OpenAcc aenv (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv (Array (sh, Int) e)
forall aenv e (acc :: * -> * -> *) sh.
Direction
-> Fun aenv (e -> e -> e)
-> Maybe (Exp aenv e)
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array (sh, Int) e)
Scan Direction
d (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall s. Exp aenv s -> Exp aenv s
cvtE (Exp aenv e -> Exp aenv e)
-> Maybe (Exp aenv e) -> Maybe (Exp aenv e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Exp aenv e)
z) (OpenAcc aenv (Array (sh, Int) e)
-> OpenAcc aenv (Array (sh, Int) e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array (sh, Int) e)
a)
        Scan' Direction
d Fun aenv (e -> e -> e)
f Exp aenv e
z OpenAcc aenv (Array (sh, Int) e)
a           -> Direction
-> Fun aenv (e -> e -> e)
-> Exp aenv e
-> OpenAcc aenv (Array (sh, Int) e)
-> PreOpenAcc OpenAcc aenv (Array (sh, Int) e, Array sh e)
forall aenv e (acc :: * -> * -> *) sh.
Direction
-> Fun aenv (e -> e -> e)
-> Exp aenv e
-> acc aenv (Array (sh, Int) e)
-> PreOpenAcc acc aenv (Array (sh, Int) e, Array sh e)
Scan' Direction
d (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (e -> e -> e)
f) (Exp aenv e -> Exp aenv e
forall s. Exp aenv s -> Exp aenv s
cvtE Exp aenv e
z) (OpenAcc aenv (Array (sh, Int) e)
-> OpenAcc aenv (Array (sh, Int) e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array (sh, Int) e)
a)
        Permute Fun aenv (e -> e -> e)
f OpenAcc aenv (Array sh' e)
d Fun aenv (sh -> PrimMaybe sh')
p OpenAcc aenv (Array sh e)
a         -> Fun aenv (e -> e -> e)
-> OpenAcc aenv (Array sh' e)
-> Fun aenv (sh -> PrimMaybe sh')
-> OpenAcc aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv (Array sh' e)
forall aenv e (acc :: * -> * -> *) sh' sh.
Fun aenv (e -> e -> e)
-> acc aenv (Array sh' e)
-> Fun aenv (sh -> PrimMaybe sh')
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh' e)
Permute (Fun aenv (e -> e -> e) -> Fun aenv (e -> e -> e)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (e -> e -> e)
f) (OpenAcc aenv (Array sh' e) -> OpenAcc aenv (Array sh' e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh' e)
d) (Fun aenv (sh -> PrimMaybe sh') -> Fun aenv (sh -> PrimMaybe sh')
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (sh -> PrimMaybe sh')
p) (OpenAcc aenv (Array sh e) -> OpenAcc aenv (Array sh e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh e)
a)
        Stencil StencilR sh e stencil
s TypeR e'
t Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
x OpenAcc aenv (Array sh e)
a       -> StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> OpenAcc aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv (Array sh e')
forall sh e stencil e' aenv (acc :: * -> * -> *).
StencilR sh e stencil
-> TypeR e'
-> Fun aenv (stencil -> e')
-> Boundary aenv (Array sh e)
-> acc aenv (Array sh e)
-> PreOpenAcc acc aenv (Array sh e')
Stencil StencilR sh e stencil
s TypeR e'
t (Fun aenv (stencil -> e') -> Fun aenv (stencil -> e')
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (stencil -> e')
f) (Boundary aenv (Array sh e) -> Boundary aenv (Array sh e)
forall s. Boundary aenv s -> Boundary aenv s
cvtB Boundary aenv (Array sh e)
x) (OpenAcc aenv (Array sh e) -> OpenAcc aenv (Array sh e)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh e)
a)
        Stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a1)
x OpenAcc aenv (Array sh a1)
a Boundary aenv (Array sh b)
y OpenAcc aenv (Array sh b)
b
                                -> StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> OpenAcc aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> OpenAcc aenv (Array sh b)
-> PreOpenAcc OpenAcc aenv (Array sh c)
forall sh a1 stencil1 b stencil2 c aenv (acc :: * -> * -> *).
StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> Fun aenv (stencil1 -> stencil2 -> c)
-> Boundary aenv (Array sh a1)
-> acc aenv (Array sh a1)
-> Boundary aenv (Array sh b)
-> acc aenv (Array sh b)
-> PreOpenAcc acc aenv (Array sh c)
Stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
t (Fun aenv (stencil1 -> stencil2 -> c)
-> Fun aenv (stencil1 -> stencil2 -> c)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (stencil1 -> stencil2 -> c)
f) (Boundary aenv (Array sh a1) -> Boundary aenv (Array sh a1)
forall s. Boundary aenv s -> Boundary aenv s
cvtB Boundary aenv (Array sh a1)
x) (OpenAcc aenv (Array sh a1) -> OpenAcc aenv (Array sh a1)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh a1)
a) (Boundary aenv (Array sh b) -> Boundary aenv (Array sh b)
forall s. Boundary aenv s -> Boundary aenv s
cvtB Boundary aenv (Array sh b)
y) (OpenAcc aenv (Array sh b) -> OpenAcc aenv (Array sh b)
forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA OpenAcc aenv (Array sh b)
b)
        -- Collect seq             -> Collect (cvtSeq seq)

      where
        cvtA :: OpenAcc aenv s -> OpenAcc aenv s
        cvtA :: forall s. OpenAcc aenv s -> OpenAcc aenv s
cvtA = (PreOpenAcc OpenAcc aenv s -> PreOpenAcc OpenAcc aenv s)
-> OpenAcc aenv s -> OpenAcc aenv s
forall aenv a b.
(PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b)
-> OpenAcc aenv a -> OpenAcc aenv b
kmap (Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv s
-> PreOpenAcc OpenAcc aenv s
forall aenv sh e a.
HasCallStack =>
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
replaceA Exp aenv sh
sh' Fun aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar)

        cvtE :: Exp aenv s -> Exp aenv s
        cvtE :: forall s. Exp aenv s -> Exp aenv s
cvtE = Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp () aenv s
-> OpenExp () aenv s
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenExp env aenv t
-> OpenExp env aenv t
replaceE Exp aenv sh
sh' Fun aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar

        cvtF :: Fun aenv s -> Fun aenv s
        cvtF :: forall s. Fun aenv s -> Fun aenv s
cvtF = Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun () aenv s
-> OpenFun () aenv s
forall env aenv sh e t.
HasCallStack =>
OpenExp env aenv sh
-> OpenFun env aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> OpenFun env aenv t
-> OpenFun env aenv t
replaceF Exp aenv sh
sh' Fun aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar

        cvtB :: Boundary aenv s -> Boundary aenv s
        cvtB :: forall s. Boundary aenv s -> Boundary aenv s
cvtB Boundary aenv s
Clamp        = Boundary aenv s
forall aenv t. Boundary aenv t
Clamp
        cvtB Boundary aenv s
Mirror       = Boundary aenv s
forall aenv t. Boundary aenv t
Mirror
        cvtB Boundary aenv s
Wrap         = Boundary aenv s
forall aenv t. Boundary aenv t
Wrap
        cvtB (Constant e
c) = e -> Boundary aenv (Array sh e)
forall e aenv sh. e -> Boundary aenv (Array sh e)
Constant e
c
        cvtB (Function Fun aenv (sh -> e)
f) = Fun aenv (sh -> e) -> Boundary aenv (Array sh e)
forall aenv sh e. Fun aenv (sh -> e) -> Boundary aenv (Array sh e)
Function (Fun aenv (sh -> e) -> Fun aenv (sh -> e)
forall s. Fun aenv s -> Fun aenv s
cvtF Fun aenv (sh -> e)
f)

        cvtAF :: HasCallStack => PreOpenAfun OpenAcc aenv s -> PreOpenAfun OpenAcc aenv s
        cvtAF :: forall s.
HasCallStack =>
PreOpenAfun OpenAcc aenv s -> PreOpenAfun OpenAcc aenv s
cvtAF = Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAfun OpenAcc aenv s
-> PreOpenAfun OpenAcc aenv s
forall aenv a.
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAfun OpenAcc aenv a
-> PreOpenAfun OpenAcc aenv a
cvt Exp aenv sh
sh' Fun aenv (sh -> e)
f' ArrayVar aenv (Array sh e)
avar
          where
            cvt :: forall aenv a.
                   Exp aenv sh -> Fun aenv (sh -> e) -> ArrayVar aenv (Array sh e)
                -> PreOpenAfun OpenAcc aenv a
                -> PreOpenAfun OpenAcc aenv a
            cvt :: forall aenv a.
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAfun OpenAcc aenv a
-> PreOpenAfun OpenAcc aenv a
cvt Exp aenv sh
sh'' Fun aenv (sh -> e)
f'' ArrayVar aenv (Array sh e)
avar' (Abody OpenAcc aenv a
a) = OpenAcc aenv a -> PreOpenAfun OpenAcc aenv a
forall (acc :: * -> * -> *) aenv t.
acc aenv t -> PreOpenAfun acc aenv t
Abody (OpenAcc aenv a -> PreOpenAfun OpenAcc aenv a)
-> OpenAcc aenv a -> PreOpenAfun OpenAcc aenv a
forall a b. (a -> b) -> a -> b
$ (PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv a)
-> OpenAcc aenv a -> OpenAcc aenv a
forall aenv a b.
(PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b)
-> OpenAcc aenv a -> OpenAcc aenv b
kmap (Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
forall aenv sh e a.
HasCallStack =>
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAcc OpenAcc aenv a
-> PreOpenAcc OpenAcc aenv a
replaceA Exp aenv sh
sh'' Fun aenv (sh -> e)
f'' ArrayVar aenv (Array sh e)
avar') OpenAcc aenv a
a
            cvt Exp aenv sh
sh'' Fun aenv (sh -> e)
f'' ArrayVar aenv (Array sh e)
avar' (Alam ALeftHandSide a aenv aenv'
lhs (PreOpenAfun OpenAcc aenv' t1
af :: PreOpenAfun OpenAcc aenv1 b)) =
              ALeftHandSide a aenv aenv'
-> PreOpenAfun OpenAcc aenv' t1
-> PreOpenAfun OpenAcc aenv (a -> t1)
forall a aenv aenv' (acc :: * -> * -> *) t1.
ALeftHandSide a aenv aenv'
-> PreOpenAfun acc aenv' t1 -> PreOpenAfun acc aenv (a -> t1)
Alam ALeftHandSide a aenv aenv'
lhs (PreOpenAfun OpenAcc aenv' t1
 -> PreOpenAfun OpenAcc aenv (a -> t1))
-> PreOpenAfun OpenAcc aenv' t1
-> PreOpenAfun OpenAcc aenv (a -> t1)
forall a b. (a -> b) -> a -> b
$ Exp aenv' sh
-> Fun aenv' (sh -> e)
-> ArrayVar aenv' (Array sh e)
-> PreOpenAfun OpenAcc aenv' t1
-> PreOpenAfun OpenAcc aenv' t1
forall aenv a.
Exp aenv sh
-> Fun aenv (sh -> e)
-> ArrayVar aenv (Array sh e)
-> PreOpenAfun OpenAcc aenv a
-> PreOpenAfun OpenAcc aenv a
cvt ((aenv :> aenv')
-> forall t. OpenExp () aenv t -> OpenExp () aenv' t
forall env env'.
(env :> env') -> forall t. OpenExp () env t -> OpenExp () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken aenv :> aenv'
w Exp aenv sh
sh'')
                ((aenv :> aenv')
-> forall t. OpenFun () aenv t -> OpenFun () aenv' t
forall env env'.
(env :> env') -> forall t. OpenFun () env t -> OpenFun () env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken aenv :> aenv'
w Fun aenv (sh -> e)
f'')
                ((aenv :> aenv')
-> forall t. Var ArrayR aenv t -> Var ArrayR aenv' t
forall env env'.
(env :> env') -> forall t. Var ArrayR env t -> Var ArrayR env' t
forall (f :: * -> * -> *) env env'.
Sink f =>
(env :> env') -> forall t. f env t -> f env' t
weaken aenv :> aenv'
w ArrayVar aenv (Array sh e)
avar')
                PreOpenAfun OpenAcc aenv' t1
af
              where
                w :: aenv :> aenv1
                w :: aenv :> aenv'
w = ALeftHandSide a aenv aenv' -> aenv :> aenv'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ALeftHandSide a aenv aenv'
lhs

-- Do not fuse bindings of multiple variables
aletD' EmbedAcc OpenAcc
_ ElimAcc OpenAcc
_ LeftHandSide ArrayR arrs aenv aenv'
lhs (Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
cc1) (Embed Extend ArrayR OpenAcc aenv' aenv'
env0 Cunctation r aenv' brrs
cc0)
  = Text -> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a. Text -> a -> a
Stats.ruleFired Text
"aletD/bind"
  (Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs)
-> Embed OpenAcc aenv brrs -> Embed OpenAcc aenv brrs
forall a b. (a -> b) -> a -> b
$ Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' brrs -> Embed OpenAcc aenv brrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed (Extend ArrayR OpenAcc aenv aenv
-> LeftHandSide ArrayR arrs aenv aenv'
-> OpenAcc aenv arrs
-> Extend ArrayR OpenAcc aenv aenv'
forall (s :: * -> *) (f :: * -> * -> *) env env'1 t env'.
Extend s f env env'1
-> LeftHandSide s t env'1 env' -> f env'1 t -> Extend s f env env'
PushEnv Extend ArrayR OpenAcc aenv aenv
forall (s :: * -> *) (f :: * -> * -> *) env. Extend s f env env
BaseEnv LeftHandSide ArrayR arrs aenv aenv'
lhs (Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (Extend ArrayR OpenAcc aenv aenv'
-> Cunctation r aenv' arrs -> Embed OpenAcc aenv arrs
forall (acc :: * -> * -> *) aenv env' e a.
Extend ArrayR acc aenv env'
-> Cunctation e env' a -> Embed acc aenv a
Embed Extend ArrayR OpenAcc aenv aenv'
env1 Cunctation r aenv' arrs
cc1)) Extend ArrayR OpenAcc aenv aenv'
-> Extend ArrayR OpenAcc aenv' aenv'
-> Extend ArrayR OpenAcc aenv aenv'
forall (s :: * -> *) (acc :: * -> * -> *) env env' env''.
Extend s acc env env'
-> Extend s acc env' env'' -> Extend s acc env env''
`append` Extend ArrayR OpenAcc aenv' aenv'
env0) Cunctation r aenv' brrs
cc0

{--
        cvtSeq :: PreOpenSeq acc aenv senv s -> PreOpenSeq acc aenv senv s
        cvtSeq s =
          case s of
            Producer p s' ->
              Producer
                (case p of
                   StreamIn arrs        -> StreamIn arrs
                   ToSeq slix sh a      -> ToSeq slix sh (cvtA a)
                   MapSeq f x           -> MapSeq (cvtAF f) x
                   ChunkedMapSeq f x    -> ChunkedMapSeq (cvtAF f) x
                   ZipWithSeq f x y     -> ZipWithSeq (cvtAF f) x y
                   ScanSeq f e x        -> ScanSeq (cvtF f) (cvtE e) x)
                (cvtSeq s')
            Consumer c ->
              Consumer (cvtC c)
            Reify ix -> Reify ix

        cvtC :: Consumer acc aenv senv s -> Consumer acc aenv senv s
        cvtC c =
          case c of
            FoldSeq f e x        -> FoldSeq (cvtF f) (cvtE e) x
            FoldSeqFlatten f a x -> FoldSeqFlatten (cvtAF f) (cvtA a) x
            Stuple t             -> Stuple (cvtCT t)

        cvtCT :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv senv) t
        cvtCT NilAtup        = NilAtup
        cvtCT (SnocAtup t c) = cvtCT t `SnocAtup` cvtC c
--}


-- Array conditionals, in particular eliminate branches when the predicate
-- reduces to a known constant.
--
-- Note that we take the raw unprocessed terms as input. If instead we had the
-- terms for each branch in the delayed representation, this would require that
-- each term has been sunk into a common environment, which implies the
-- conditional has been pushed underneath the intersection of bound terms for
-- both branches. This would result in redundant work processing the bindings
-- for the branch not taken.
--
acondD :: HasCallStack
       => MatchAcc OpenAcc
       -> EmbedAcc OpenAcc
       -> Exp              aenv PrimBool
       ->          OpenAcc aenv arrs
       ->          OpenAcc aenv arrs
       -> Embed    OpenAcc aenv arrs
acondD :: forall aenv arrs.
HasCallStack =>
MatchAcc OpenAcc
-> EmbedAcc OpenAcc
-> Exp aenv PrimBool
-> OpenAcc aenv arrs
-> OpenAcc aenv arrs
-> Embed OpenAcc aenv arrs
acondD MatchAcc OpenAcc
matchAcc EmbedAcc OpenAcc
embedAcc Exp aenv PrimBool
p OpenAcc aenv arrs
t OpenAcc aenv arrs
e
  | Const ScalarType PrimBool
_ PrimBool
1 <- Exp aenv PrimBool
p            = Text -> Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a. Text -> a -> a
Stats.knownBranch Text
"True"      (Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv arrs
t
  | Const ScalarType PrimBool
_ PrimBool
0 <- Exp aenv PrimBool
p            = Text -> Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a. Text -> a -> a
Stats.knownBranch Text
"False"     (Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv arrs
e
  | Just arrs :~: arrs
Refl <- OpenAcc aenv arrs -> OpenAcc aenv arrs -> Maybe (arrs :~: arrs)
MatchAcc OpenAcc
matchAcc OpenAcc aenv arrs
t OpenAcc aenv arrs
e = Text -> Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a. Text -> a -> a
Stats.knownBranch Text
"redundant" (Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv arrs
e
  | Bool
otherwise                 = PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall aenv a.
HasCallStack =>
PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a
done (PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs)
-> PreOpenAcc OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
forall a b. (a -> b) -> a -> b
$ Exp aenv PrimBool
-> OpenAcc aenv arrs
-> OpenAcc aenv arrs
-> PreOpenAcc OpenAcc aenv arrs
forall aenv (acc :: * -> * -> *) a.
Exp aenv PrimBool
-> acc aenv a -> acc aenv a -> PreOpenAcc acc aenv a
Acond Exp aenv PrimBool
p (Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv arrs
t))
                                               (Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
forall aenv arrs.
HasCallStack =>
Embed OpenAcc aenv arrs -> OpenAcc aenv arrs
computeAcc (OpenAcc aenv arrs -> Embed OpenAcc aenv arrs
EmbedAcc OpenAcc
embedAcc OpenAcc aenv arrs
e))


-- Scalar expressions
-- ------------------

identity :: TypeR a -> OpenFun env aenv (a -> a)
identity :: forall a env aenv. TypeR a -> OpenFun env aenv (a -> a)
identity TypeR a
t
  | DeclareVars LeftHandSide ScalarType a env env'
lhs env :> env'
_ forall env''. (env' :> env'') -> Vars ScalarType env'' a
value <- TypeR a -> DeclareVars ScalarType a env
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars TypeR a
t
  = LeftHandSide ScalarType a env env'
-> OpenFun env' aenv a -> OpenFun env aenv (a -> a)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam LeftHandSide ScalarType a env env'
lhs (OpenFun env' aenv a -> OpenFun env aenv (a -> a))
-> OpenFun env' aenv a -> OpenFun env aenv (a -> a)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv a -> OpenFun env' aenv a
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv a -> OpenFun env' aenv a)
-> OpenExp env' aenv a -> OpenFun env' aenv a
forall a b. (a -> b) -> a -> b
$ ExpVars env' a -> OpenExp env' aenv a
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars (ExpVars env' a -> OpenExp env' aenv a)
-> ExpVars env' a -> OpenExp env' aenv a
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ExpVars env' a
forall env''. (env' :> env'') -> Vars ScalarType env'' a
value env' :> env'
forall env. env :> env
weakenId

toIndex :: ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (sh -> Int)
toIndex :: forall sh env aenv.
ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (sh -> Int)
toIndex ShapeR sh
shR OpenExp env aenv sh
sh
  | DeclareVars LeftHandSide ScalarType sh env env'
lhs env :> env'
k forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value <- TupR ScalarType sh -> DeclareVars ScalarType sh env
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (TupR ScalarType sh -> DeclareVars ScalarType sh env)
-> TupR ScalarType sh -> DeclareVars ScalarType sh env
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> TupR ScalarType sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shR
  = LeftHandSide ScalarType sh env env'
-> OpenFun env' aenv Int -> OpenFun env aenv (sh -> Int)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam LeftHandSide ScalarType sh env env'
lhs (OpenFun env' aenv Int -> OpenFun env aenv (sh -> Int))
-> OpenFun env' aenv Int -> OpenFun env aenv (sh -> Int)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv Int -> OpenFun env' aenv Int
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv Int -> OpenFun env' aenv Int)
-> OpenExp env' aenv Int -> OpenFun env' aenv Int
forall a b. (a -> b) -> a -> b
$ ShapeR sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv Int
forall sh env aenv.
ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv Int
ToIndex ShapeR sh
shR ((env :> env') -> OpenExp env aenv sh -> OpenExp env' aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env :> env'
k OpenExp env aenv sh
sh) (OpenExp env' aenv sh -> OpenExp env' aenv Int)
-> OpenExp env' aenv sh -> OpenExp env' aenv Int
forall a b. (a -> b) -> a -> b
$ ExpVars env' sh -> OpenExp env' aenv sh
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars (ExpVars env' sh -> OpenExp env' aenv sh)
-> ExpVars env' sh -> OpenExp env' aenv sh
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ExpVars env' sh
forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value env' :> env'
forall env. env :> env
weakenId

fromIndex :: ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (Int -> sh)
fromIndex :: forall sh env aenv.
ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (Int -> sh)
fromIndex ShapeR sh
shR OpenExp env aenv sh
sh
  = ELeftHandSide Int env (env, Int)
-> OpenFun (env, Int) aenv sh -> OpenFun env aenv (Int -> sh)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam (ScalarType Int -> ELeftHandSide Int env (env, Int)
forall (s :: * -> *) v env. s v -> LeftHandSide s v env (env, v)
LeftHandSideSingle ScalarType Int
scalarTypeInt)
  (OpenFun (env, Int) aenv sh -> OpenFun env aenv (Int -> sh))
-> OpenFun (env, Int) aenv sh -> OpenFun env aenv (Int -> sh)
forall a b. (a -> b) -> a -> b
$ OpenExp (env, Int) aenv sh -> OpenFun (env, Int) aenv sh
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body
  (OpenExp (env, Int) aenv sh -> OpenFun (env, Int) aenv sh)
-> OpenExp (env, Int) aenv sh -> OpenFun (env, Int) aenv sh
forall a b. (a -> b) -> a -> b
$ ShapeR sh
-> OpenExp (env, Int) aenv sh
-> OpenExp (env, Int) aenv Int
-> OpenExp (env, Int) aenv sh
forall t env aenv.
ShapeR t
-> OpenExp env aenv t -> OpenExp env aenv Int -> OpenExp env aenv t
FromIndex ShapeR sh
shR ((env :> (env, Int))
-> OpenExp env aenv sh -> OpenExp (env, Int) aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE ((env :> env) -> env :> (env, Int)
forall env env' t. (env :> env') -> env :> (env', t)
weakenSucc' env :> env
forall env. env :> env
weakenId) OpenExp env aenv sh
sh)
  (OpenExp (env, Int) aenv Int -> OpenExp (env, Int) aenv sh)
-> OpenExp (env, Int) aenv Int -> OpenExp (env, Int) aenv sh
forall a b. (a -> b) -> a -> b
$ ExpVar (env, Int) Int -> OpenExp (env, Int) aenv Int
forall env t aenv. ExpVar env t -> OpenExp env aenv t
Evar
  (ExpVar (env, Int) Int -> OpenExp (env, Int) aenv Int)
-> ExpVar (env, Int) Int -> OpenExp (env, Int) aenv Int
forall a b. (a -> b) -> a -> b
$ ScalarType Int -> Idx (env, Int) Int -> ExpVar (env, Int) Int
forall (s :: * -> *) env t. s t -> Idx env t -> Var s env t
Var ScalarType Int
scalarTypeInt Idx (env, Int) Int
forall envt t env. (envt ~ (env, t)) => Idx envt t
ZeroIdx

intersect :: ShapeR sh -> OpenExp env aenv sh -> OpenExp env aenv sh -> OpenExp env aenv sh
intersect :: forall sh env aenv.
ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
intersect = (forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
forall aenv sh env.
(forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
mkShapeBinary OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
forall {t} {env} {aenv}.
IsSingle t =>
OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t
f
  where
    f :: OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t
f OpenExp env aenv t
a OpenExp env aenv t
b = PrimFun ((t, t) -> t)
-> OpenExp env aenv (t, t) -> OpenExp env aenv t
forall a t env aenv.
PrimFun (a -> t) -> OpenExp env aenv a -> OpenExp env aenv t
PrimApp (SingleType t -> PrimFun ((t, t) -> t)
forall a. SingleType a -> PrimFun ((a, a) -> a)
PrimMin SingleType t
forall a. IsSingle a => SingleType a
singleType) (OpenExp env aenv (t, t) -> OpenExp env aenv t)
-> OpenExp env aenv (t, t) -> OpenExp env aenv t
forall a b. (a -> b) -> a -> b
$ OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv (t, t)
forall env aenv t1 t2.
OpenExp env aenv t1
-> OpenExp env aenv t2 -> OpenExp env aenv (t1, t2)
Pair OpenExp env aenv t
a OpenExp env aenv t
b

-- union :: ShapeR sh -> OpenExp env aenv sh -> OpenExp env aenv sh -> OpenExp env aenv sh
-- union = mkShapeBinary f
--   where
--     f a b = PrimApp (PrimMax singleType) $ Pair a b

mkShapeBinary
    :: (forall env'. OpenExp env' aenv Int -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
    -> ShapeR sh
    -> OpenExp env aenv sh
    -> OpenExp env aenv sh
    -> OpenExp env aenv sh
mkShapeBinary :: forall aenv sh env.
(forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
mkShapeBinary forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
_ ShapeR sh
ShapeRz OpenExp env aenv sh
_ OpenExp env aenv sh
_ = OpenExp env aenv sh
OpenExp env aenv ()
forall env aenv. OpenExp env aenv ()
Nil
mkShapeBinary forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f (ShapeRsnoc ShapeR sh1
shR) (Pair OpenExp env aenv t1
as OpenExp env aenv t2
a) (Pair OpenExp env aenv t1
bs OpenExp env aenv t2
b) = (forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh1
-> OpenExp env aenv sh1
-> OpenExp env aenv sh1
-> OpenExp env aenv sh1
forall aenv sh env.
(forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
mkShapeBinary OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh1
shR OpenExp env aenv sh1
OpenExp env aenv t1
as OpenExp env aenv sh1
OpenExp env aenv t1
bs OpenExp env aenv sh1
-> OpenExp env aenv Int -> OpenExp env aenv (sh1, Int)
forall env aenv t1 t2.
OpenExp env aenv t1
-> OpenExp env aenv t2 -> OpenExp env aenv (t1, t2)
`Pair` OpenExp env aenv Int
-> OpenExp env aenv Int -> OpenExp env aenv Int
forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f OpenExp env aenv t2
OpenExp env aenv Int
a OpenExp env aenv t2
OpenExp env aenv Int
b
mkShapeBinary forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR (Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
bnd OpenExp env' aenv sh
a) OpenExp env aenv sh
b = ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv sh
-> OpenExp env aenv sh
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
bnd (OpenExp env' aenv sh -> OpenExp env aenv sh)
-> OpenExp env' aenv sh -> OpenExp env aenv sh
forall a b. (a -> b) -> a -> b
$ (forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
forall aenv sh env.
(forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
mkShapeBinary OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR OpenExp env' aenv sh
a ((env :> env') -> OpenExp env aenv sh -> OpenExp env' aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE (ELeftHandSide bnd_t env env' -> env :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ELeftHandSide bnd_t env env'
lhs) OpenExp env aenv sh
b)
mkShapeBinary forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR OpenExp env aenv sh
a (Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
bnd OpenExp env' aenv sh
b) = ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv sh
-> OpenExp env aenv sh
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
bnd (OpenExp env' aenv sh -> OpenExp env aenv sh)
-> OpenExp env' aenv sh -> OpenExp env aenv sh
forall a b. (a -> b) -> a -> b
$ (forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
forall aenv sh env.
(forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
mkShapeBinary OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR ((env :> env') -> OpenExp env aenv sh -> OpenExp env' aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE (ELeftHandSide bnd_t env env' -> env :> env'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ELeftHandSide bnd_t env env'
lhs) OpenExp env aenv sh
a) OpenExp env' aenv sh
b
mkShapeBinary forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR OpenExp env aenv sh
a b :: OpenExp env aenv sh
b@Pair{} -- `a` is not Pair
  | DeclareVars LeftHandSide ScalarType sh env env'
lhs env :> env'
k forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value <- TupR ScalarType sh -> DeclareVars ScalarType sh env
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (TupR ScalarType sh -> DeclareVars ScalarType sh env)
-> TupR ScalarType sh -> DeclareVars ScalarType sh env
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> TupR ScalarType sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shR
  = LeftHandSide ScalarType sh env env'
-> OpenExp env aenv sh
-> OpenExp env' aenv sh
-> OpenExp env aenv sh
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let LeftHandSide ScalarType sh env env'
lhs OpenExp env aenv sh
a (OpenExp env' aenv sh -> OpenExp env aenv sh)
-> OpenExp env' aenv sh -> OpenExp env aenv sh
forall a b. (a -> b) -> a -> b
$ (forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
forall aenv sh env.
(forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
mkShapeBinary OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR (ExpVars env' sh -> OpenExp env' aenv sh
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars (ExpVars env' sh -> OpenExp env' aenv sh)
-> ExpVars env' sh -> OpenExp env' aenv sh
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ExpVars env' sh
forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value env' :> env'
forall env. env :> env
weakenId) ((env :> env') -> OpenExp env aenv sh -> OpenExp env' aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env :> env'
k OpenExp env aenv sh
b)
mkShapeBinary forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR OpenExp env aenv sh
a OpenExp env aenv sh
b -- `b` is not a Pair
  | DeclareVars LeftHandSide ScalarType sh env env'
lhs env :> env'
k forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value <- TupR ScalarType sh -> DeclareVars ScalarType sh env
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (TupR ScalarType sh -> DeclareVars ScalarType sh env)
-> TupR ScalarType sh -> DeclareVars ScalarType sh env
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> TupR ScalarType sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shR
  = LeftHandSide ScalarType sh env env'
-> OpenExp env aenv sh
-> OpenExp env' aenv sh
-> OpenExp env aenv sh
forall bnd_t env env' aenv t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
Let LeftHandSide ScalarType sh env env'
lhs OpenExp env aenv sh
b (OpenExp env' aenv sh -> OpenExp env aenv sh)
-> OpenExp env' aenv sh -> OpenExp env aenv sh
forall a b. (a -> b) -> a -> b
$ (forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
-> OpenExp env' aenv sh
forall aenv sh env.
(forall env'.
 OpenExp env' aenv Int
 -> OpenExp env' aenv Int -> OpenExp env' aenv Int)
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
mkShapeBinary OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
forall env'.
OpenExp env' aenv Int
-> OpenExp env' aenv Int -> OpenExp env' aenv Int
f ShapeR sh
shR ((env :> env') -> OpenExp env aenv sh -> OpenExp env' aenv sh
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE env :> env'
k OpenExp env aenv sh
a) (ExpVars env' sh -> OpenExp env' aenv sh
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars (ExpVars env' sh -> OpenExp env' aenv sh)
-> ExpVars env' sh -> OpenExp env' aenv sh
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ExpVars env' sh
forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value env' :> env'
forall env. env :> env
weakenId)

reindex :: ShapeR sh'
        -> OpenExp env aenv sh'
        -> ShapeR sh
        -> OpenExp env aenv sh
        -> OpenFun env aenv (sh -> sh')
reindex :: forall sh' env aenv sh.
ShapeR sh'
-> OpenExp env aenv sh'
-> ShapeR sh
-> OpenExp env aenv sh
-> OpenFun env aenv (sh -> sh')
reindex ShapeR sh'
shR' OpenExp env aenv sh'
sh' ShapeR sh
shR OpenExp env aenv sh
sh
  | Just sh :~: sh'
Refl <- OpenExp env aenv sh -> OpenExp env aenv sh' -> Maybe (sh :~: sh')
forall env aenv s t.
OpenExp env aenv s -> OpenExp env aenv t -> Maybe (s :~: t)
matchOpenExp OpenExp env aenv sh
sh OpenExp env aenv sh'
sh' = TypeR sh -> OpenFun env aenv (sh -> sh)
forall a env aenv. TypeR a -> OpenFun env aenv (a -> a)
identity (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh'
ShapeR sh
shR')
  | Bool
otherwise                        = ShapeR sh' -> OpenExp env aenv sh' -> OpenFun env aenv (Int -> sh')
forall sh env aenv.
ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (Int -> sh)
fromIndex ShapeR sh'
shR' OpenExp env aenv sh'
sh' OpenFun env aenv (Int -> sh')
-> OpenFun env aenv (sh -> Int) -> OpenFun env aenv (sh -> sh')
forall env aenv b c a.
HasCallStack =>
OpenFun env aenv (b -> c)
-> OpenFun env aenv (a -> b) -> OpenFun env aenv (a -> c)
`compose` ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (sh -> Int)
forall sh env aenv.
ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (sh -> Int)
toIndex ShapeR sh
shR OpenExp env aenv sh
sh

extend :: SliceIndex slix sl co sh
       -> Exp aenv slix
       -> Fun aenv (sh -> sl)
extend :: forall slix sl co sh aenv.
SliceIndex slix sl co sh -> Exp aenv slix -> Fun aenv (sh -> sl)
extend SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix
  | DeclareVars LeftHandSide ScalarType sh () env'
lhs () :> env'
k forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value <- TupR ScalarType sh -> DeclareVars ScalarType sh ()
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (TupR ScalarType sh -> DeclareVars ScalarType sh ())
-> TupR ScalarType sh -> DeclareVars ScalarType sh ()
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> TupR ScalarType sh
forall sh. ShapeR sh -> TypeR sh
shapeType (ShapeR sh -> TupR ScalarType sh)
-> ShapeR sh -> TupR ScalarType sh
forall a b. (a -> b) -> a -> b
$ SliceIndex slix sl co sh -> ShapeR sh
forall slix sl co dim. SliceIndex slix sl co dim -> ShapeR dim
sliceDomainR SliceIndex slix sl co sh
sliceIndex
  = LeftHandSide ScalarType sh () env'
-> OpenFun env' aenv sl -> OpenFun () aenv (sh -> sl)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam LeftHandSide ScalarType sh () env'
lhs (OpenFun env' aenv sl -> OpenFun () aenv (sh -> sl))
-> OpenFun env' aenv sl -> OpenFun () aenv (sh -> sl)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv sl -> OpenFun env' aenv sl
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv sl -> OpenFun env' aenv sl)
-> OpenExp env' aenv sl -> OpenFun env' aenv sl
forall a b. (a -> b) -> a -> b
$ SliceIndex slix sl co sh
-> OpenExp env' aenv slix
-> OpenExp env' aenv sh
-> OpenExp env' aenv sl
forall slix t co sh env aenv.
SliceIndex slix t co sh
-> OpenExp env aenv slix
-> OpenExp env aenv sh
-> OpenExp env aenv t
IndexSlice SliceIndex slix sl co sh
sliceIndex ((() :> env') -> Exp aenv slix -> OpenExp env' aenv slix
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE () :> env'
k Exp aenv slix
slix) (OpenExp env' aenv sh -> OpenExp env' aenv sl)
-> OpenExp env' aenv sh -> OpenExp env' aenv sl
forall a b. (a -> b) -> a -> b
$ ExpVars env' sh -> OpenExp env' aenv sh
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars (ExpVars env' sh -> OpenExp env' aenv sh)
-> ExpVars env' sh -> OpenExp env' aenv sh
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ExpVars env' sh
forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value env' :> env'
forall env. env :> env
weakenId

restrict :: SliceIndex slix sl co sh
         -> Exp aenv slix
         -> Fun aenv (sl -> sh)
restrict :: forall slix sl co sh aenv.
SliceIndex slix sl co sh -> Exp aenv slix -> Fun aenv (sl -> sh)
restrict SliceIndex slix sl co sh
sliceIndex Exp aenv slix
slix
  | DeclareVars LeftHandSide ScalarType sl () env'
lhs () :> env'
k forall env''. (env' :> env'') -> Vars ScalarType env'' sl
value <- TupR ScalarType sl -> DeclareVars ScalarType sl ()
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (TupR ScalarType sl -> DeclareVars ScalarType sl ())
-> TupR ScalarType sl -> DeclareVars ScalarType sl ()
forall a b. (a -> b) -> a -> b
$ ShapeR sl -> TupR ScalarType sl
forall sh. ShapeR sh -> TypeR sh
shapeType (ShapeR sl -> TupR ScalarType sl)
-> ShapeR sl -> TupR ScalarType sl
forall a b. (a -> b) -> a -> b
$ SliceIndex slix sl co sh -> ShapeR sl
forall slix sl co dim. SliceIndex slix sl co dim -> ShapeR sl
sliceShapeR SliceIndex slix sl co sh
sliceIndex
  = LeftHandSide ScalarType sl () env'
-> OpenFun env' aenv sh -> OpenFun () aenv (sl -> sh)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam LeftHandSide ScalarType sl () env'
lhs (OpenFun env' aenv sh -> OpenFun () aenv (sl -> sh))
-> OpenFun env' aenv sh -> OpenFun () aenv (sl -> sh)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv sh -> OpenFun env' aenv sh
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv sh -> OpenFun env' aenv sh)
-> OpenExp env' aenv sh -> OpenFun env' aenv sh
forall a b. (a -> b) -> a -> b
$ SliceIndex slix sl co sh
-> OpenExp env' aenv slix
-> OpenExp env' aenv sl
-> OpenExp env' aenv sh
forall slix sl co t env aenv.
SliceIndex slix sl co t
-> OpenExp env aenv slix
-> OpenExp env aenv sl
-> OpenExp env aenv t
IndexFull SliceIndex slix sl co sh
sliceIndex ((() :> env') -> Exp aenv slix -> OpenExp env' aenv slix
forall env env' aenv t.
(env :> env') -> OpenExp env aenv t -> OpenExp env' aenv t
forall (f :: * -> * -> * -> *) env env' aenv t.
SinkExp f =>
(env :> env') -> f env aenv t -> f env' aenv t
weakenE () :> env'
k Exp aenv slix
slix) (OpenExp env' aenv sl -> OpenExp env' aenv sh)
-> OpenExp env' aenv sl -> OpenExp env' aenv sh
forall a b. (a -> b) -> a -> b
$ ExpVars env' sl -> OpenExp env' aenv sl
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars (ExpVars env' sl -> OpenExp env' aenv sl)
-> ExpVars env' sl -> OpenExp env' aenv sl
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> ExpVars env' sl
forall env''. (env' :> env'') -> Vars ScalarType env'' sl
value env' :> env'
forall env. env :> env
weakenId

arrayShape :: ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape :: forall aenv sh e. ArrayVar aenv (Array sh e) -> Exp aenv sh
arrayShape = Exp aenv sh -> Exp aenv sh
forall aenv' t. HasCallStack => Exp aenv' t -> Exp aenv' t
simplifyExp (Exp aenv sh -> Exp aenv sh)
-> (ArrayVar aenv (Array sh e) -> Exp aenv sh)
-> ArrayVar aenv (Array sh e)
-> Exp aenv sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ArrayVar aenv (Array sh e) -> Exp aenv sh
forall aenv t e env.
ArrayVar aenv (Array t e) -> OpenExp env aenv t
Shape

indexArray :: ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray :: forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e)
indexArray v :: ArrayVar aenv (Array sh e)
v@(Var (ArrayR ShapeR sh
shR TypeR e
_) Idx aenv (Array sh e)
_)
  | DeclareVars LeftHandSide ScalarType sh () env'
lhs () :> env'
_ forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value <- TupR ScalarType sh -> DeclareVars ScalarType sh ()
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars (TupR ScalarType sh -> DeclareVars ScalarType sh ())
-> TupR ScalarType sh -> DeclareVars ScalarType sh ()
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> TupR ScalarType sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shR
  = ELeftHandSide sh () env'
-> OpenFun env' aenv e -> OpenFun () aenv (sh -> e)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam ELeftHandSide sh () env'
LeftHandSide ScalarType sh () env'
lhs (OpenFun env' aenv e -> OpenFun () aenv (sh -> e))
-> OpenFun env' aenv e -> OpenFun () aenv (sh -> e)
forall a b. (a -> b) -> a -> b
$ OpenExp env' aenv e -> OpenFun env' aenv e
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env' aenv e -> OpenFun env' aenv e)
-> OpenExp env' aenv e -> OpenFun env' aenv e
forall a b. (a -> b) -> a -> b
$ ArrayVar aenv (Array sh e)
-> OpenExp env' aenv sh -> OpenExp env' aenv e
forall aenv dim t env.
ArrayVar aenv (Array dim t)
-> OpenExp env aenv dim -> OpenExp env aenv t
Index ArrayVar aenv (Array sh e)
v (OpenExp env' aenv sh -> OpenExp env' aenv e)
-> OpenExp env' aenv sh -> OpenExp env' aenv e
forall a b. (a -> b) -> a -> b
$ ExpVars env' sh -> OpenExp env' aenv sh
forall env t aenv. ExpVars env t -> OpenExp env aenv t
expVars (ExpVars env' sh -> OpenExp env' aenv sh)
-> ExpVars env' sh -> OpenExp env' aenv sh
forall a b. (a -> b) -> a -> b
$ (env' :> env') -> Vars ScalarType env' sh
forall env''. (env' :> env'') -> Vars ScalarType env'' sh
value env' :> env'
forall env. env :> env
weakenId

linearIndex :: ArrayVar aenv (Array sh e) -> Fun aenv (Int -> e)
linearIndex :: forall aenv sh e. ArrayVar aenv (Array sh e) -> Fun aenv (Int -> e)
linearIndex ArrayVar aenv (Array sh e)
v = ELeftHandSide Int () ((), Int)
-> OpenFun ((), Int) aenv e -> OpenFun () aenv (Int -> e)
forall a env env' aenv t1.
ELeftHandSide a env env'
-> OpenFun env' aenv t1 -> OpenFun env aenv (a -> t1)
Lam (ScalarType Int -> ELeftHandSide Int () ((), Int)
forall (s :: * -> *) v env. s v -> LeftHandSide s v env (env, v)
LeftHandSideSingle ScalarType Int
scalarTypeInt) (OpenFun ((), Int) aenv e -> OpenFun () aenv (Int -> e))
-> OpenFun ((), Int) aenv e -> OpenFun () aenv (Int -> e)
forall a b. (a -> b) -> a -> b
$ OpenExp ((), Int) aenv e -> OpenFun ((), Int) aenv e
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp ((), Int) aenv e -> OpenFun ((), Int) aenv e)
-> OpenExp ((), Int) aenv e -> OpenFun ((), Int) aenv e
forall a b. (a -> b) -> a -> b
$ ArrayVar aenv (Array sh e)
-> OpenExp ((), Int) aenv Int -> OpenExp ((), Int) aenv e
forall aenv dim t env.
ArrayVar aenv (Array dim t)
-> OpenExp env aenv Int -> OpenExp env aenv t
LinearIndex ArrayVar aenv (Array sh e)
v (OpenExp ((), Int) aenv Int -> OpenExp ((), Int) aenv e)
-> OpenExp ((), Int) aenv Int -> OpenExp ((), Int) aenv e
forall a b. (a -> b) -> a -> b
$ ExpVar ((), Int) Int -> OpenExp ((), Int) aenv Int
forall env t aenv. ExpVar env t -> OpenExp env aenv t
Evar (ExpVar ((), Int) Int -> OpenExp ((), Int) aenv Int)
-> ExpVar ((), Int) Int -> OpenExp ((), Int) aenv Int
forall a b. (a -> b) -> a -> b
$ ScalarType Int -> Idx ((), Int) Int -> ExpVar ((), Int) Int
forall (s :: * -> *) env t. s t -> Idx env t -> Var s env t
Var ScalarType Int
scalarTypeInt Idx ((), Int) Int
forall envt t env. (envt ~ (env, t)) => Idx envt t
ZeroIdx


extractOpenAcc :: ExtractAcc OpenAcc
extractOpenAcc :: ExtractAcc OpenAcc
extractOpenAcc (OpenAcc PreOpenAcc OpenAcc env t
pacc) = PreOpenAcc OpenAcc env t -> Maybe (PreOpenAcc OpenAcc env t)
forall a. a -> Maybe a
Just PreOpenAcc OpenAcc env t
pacc

extractDelayedOpenAcc :: ExtractAcc DelayedOpenAcc
extractDelayedOpenAcc :: ExtractAcc DelayedOpenAcc
extractDelayedOpenAcc (Manifest PreOpenAcc DelayedOpenAcc env t
pacc) = PreOpenAcc DelayedOpenAcc env t
-> Maybe (PreOpenAcc DelayedOpenAcc env t)
forall a. a -> Maybe a
Just PreOpenAcc DelayedOpenAcc env t
pacc
extractDelayedOpenAcc DelayedOpenAcc env t
_               = Maybe (PreOpenAcc DelayedOpenAcc env t)
forall a. Maybe a
Nothing

extractOpenArrayVars
    :: OpenAcc aenv a
    -> Maybe (ArrayVars aenv a)
extractOpenArrayVars :: forall aenv a. OpenAcc aenv a -> Maybe (ArrayVars aenv a)
extractOpenArrayVars (OpenAcc PreOpenAcc OpenAcc aenv a
pacc) =
  ExtractAcc OpenAcc
-> PreOpenAcc OpenAcc aenv a -> Maybe (ArrayVars aenv a)
forall (acc :: * -> * -> *) aenv a.
ExtractAcc acc -> PreOpenAcc acc aenv a -> Maybe (ArrayVars aenv a)
avarsOut OpenAcc env t -> Maybe (PreOpenAcc OpenAcc env t)
ExtractAcc OpenAcc
extractOpenAcc PreOpenAcc OpenAcc aenv a
pacc

extractDelayedArrayVars
    :: DelayedOpenAcc aenv a
    -> Maybe (ArrayVars aenv a)
extractDelayedArrayVars :: forall aenv a. DelayedOpenAcc aenv a -> Maybe (ArrayVars aenv a)
extractDelayedArrayVars DelayedOpenAcc aenv a
acc
  | Just PreOpenAcc DelayedOpenAcc aenv a
pacc <- DelayedOpenAcc aenv a -> Maybe (PreOpenAcc DelayedOpenAcc aenv a)
ExtractAcc DelayedOpenAcc
extractDelayedOpenAcc DelayedOpenAcc aenv a
acc = ExtractAcc DelayedOpenAcc
-> PreOpenAcc DelayedOpenAcc aenv a -> Maybe (ArrayVars aenv a)
forall (acc :: * -> * -> *) aenv a.
ExtractAcc acc -> PreOpenAcc acc aenv a -> Maybe (ArrayVars aenv a)
avarsOut DelayedOpenAcc env t -> Maybe (PreOpenAcc DelayedOpenAcc env t)
ExtractAcc DelayedOpenAcc
extractDelayedOpenAcc PreOpenAcc DelayedOpenAcc aenv a
pacc
  | Bool
otherwise                              = Maybe (ArrayVars aenv a)
forall a. Maybe a
Nothing