{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards         #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE RebindableSyntax      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}   -- pattern synonyms
-- |
-- Module      : Data.Array.Accelerate.Prelude
-- Copyright   : [2009..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Standard functions that are not part of the core set (directly represented in
-- the AST), but are instead implemented in terms of the core set.
--

module Data.Array.Accelerate.Prelude (

  -- * Element-wise operations
  indexed,
  imap,

  -- * Zipping
  zipWith3, zipWith4, zipWith5, zipWith6, zipWith7, zipWith8, zipWith9,
  izipWith, izipWith3, izipWith4, izipWith5, izipWith6, izipWith7, izipWith8, izipWith9,
  zip, zip3, zip4, zip5, zip6, zip7, zip8, zip9,

  -- * Unzipping
  unzip, unzip3, unzip4, unzip5, unzip6, unzip7, unzip8, unzip9,

  -- * Reductions
  foldAll, fold1All,
  foldSeg, fold1Seg,

  -- ** Specialised folds
  all, any, and, or, sum, product, minimum, maximum,

  -- * Scans
  prescanl, postscanl, prescanr, postscanr,

  -- ** Segmented scans
  scanlSeg, scanl'Seg, scanl1Seg, prescanlSeg, postscanlSeg,
  scanrSeg, scanr'Seg, scanr1Seg, prescanrSeg, postscanrSeg,

  -- * Shape manipulation
  flatten,

  -- * Enumeration and filling
  fill, enumFromN, enumFromStepN,

  -- * Concatenation
  (++), concatOn,

  -- * Working with predicates
  -- ** Filtering
  filter, compact,

  -- ** Scatter / Gather
  scatter, scatterIf,
  gather,  gatherIf,

  -- * Permutations
  reverse, transpose,
  reverseOn, transposeOn,

  -- * Extracting sub-vectors
  init, tail, take, drop, slit,
  initOn, tailOn, takeOn, dropOn, slitOn,

  -- * Controlling execution
  compute,

  -- * Flow control
  IfThenElse(..),

  -- ** Array-level
  (?|),

  -- ** Expression-level
  (?), match,

  -- * Scalar iteration
  iterate,

  -- * Scalar reduction
  sfoldl, -- sfoldr,

  -- * Lifting and unlifting
  Lift(..), Unlift(..),
  lift1, lift2, lift3, ilift1, ilift2, ilift3,

  -- ** Tuple construction and destruction
  fst, afst, snd, asnd, curry, uncurry,

  -- ** Index construction and destruction
  index0, index1, unindex1, index2, unindex2, index3, unindex3,

  -- * Array operations with a scalar result
  the, null, length,

  -- * Irregular data-parallelism
  expand,

  -- * Sequence operations
  -- fromSeq, fromSeqElems, fromSeqShapes, toSeqInner, toSeqOuter2, toSeqOuter3, generateSeq,

) where

import Data.Array.Accelerate.Analysis.Match
import Data.Array.Accelerate.Language
import Data.Array.Accelerate.Lift
import Data.Array.Accelerate.Pattern
import Data.Array.Accelerate.Pattern.Maybe
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Sugar.Array                            ( Arrays, Array, Scalar, Vector, Segments,  fromList )
import Data.Array.Accelerate.Sugar.Elt
import Data.Array.Accelerate.Sugar.Shape                            ( Shape, Slice, Z(..), (:.)(..), All(..), DIM1, DIM2, empty )
import Data.Array.Accelerate.Type

import Data.Array.Accelerate.Classes.Eq
import Data.Array.Accelerate.Classes.FromIntegral
import Data.Array.Accelerate.Classes.Integral
import Data.Array.Accelerate.Classes.Num
import Data.Array.Accelerate.Classes.Ord

import Data.Array.Accelerate.Data.Bits

import Control.Lens                                                 ( Lens', (&), (^.), (.~), (+~), (-~), lens, over )
import Prelude                                                      ( (.), ($), Maybe(..), const, id, flip )

import GHC.Base                                                     ( Constraint )


-- $setup
-- >>> :seti -XFlexibleContexts
-- >>> import Data.Array.Accelerate
-- >>> import Data.Array.Accelerate.Interpreter
-- >>> :{
--   let runExp :: Elt e => Exp e -> e
--       runExp e = indexArray (run (unit e)) Z
-- :}

-- Element-wise operations
-- -----------------------

-- | Pair each element with its index
--
-- >>> let xs = fromList (Z:.5) [0..] :: Vector Float
-- >>> run $ indexed (use xs)
-- Vector (Z :. 5) [(Z :. 0,0.0),(Z :. 1,1.0),(Z :. 2,2.0),(Z :. 3,3.0),(Z :. 4,4.0)]
--
-- >>> let mat = fromList (Z:.3:.4) [0..] :: Matrix Float
-- >>> run $ indexed (use mat)
-- Matrix (Z :. 3 :. 4)
--   [ (Z :. 0 :. 0,0.0), (Z :. 0 :. 1,1.0),  (Z :. 0 :. 2,2.0),  (Z :. 0 :. 3,3.0),
--     (Z :. 1 :. 0,4.0), (Z :. 1 :. 1,5.0),  (Z :. 1 :. 2,6.0),  (Z :. 1 :. 3,7.0),
--     (Z :. 2 :. 0,8.0), (Z :. 2 :. 1,9.0), (Z :. 2 :. 2,10.0), (Z :. 2 :. 3,11.0)]
--
indexed :: (Shape sh, Elt a) => Acc (Array sh a) -> Acc (Array sh (sh, a))
indexed :: Acc (Array sh a) -> Acc (Array sh (sh, a))
indexed = (Exp sh -> Exp a -> Exp (sh, a))
-> Acc (Array sh a) -> Acc (Array sh (sh, a))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp sh -> Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
imap Exp sh -> Exp a -> Exp (sh, a)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2

-- | Apply a function to every element of an array and its index
--
imap :: (Shape sh, Elt a, Elt b)
     => (Exp sh -> Exp a -> Exp b)
     -> Acc (Array sh a)
     -> Acc (Array sh b)
imap :: (Exp sh -> Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
imap Exp sh -> Exp a -> Exp b
f Acc (Array sh a)
xs = (Exp sh -> Exp a -> Exp b)
-> Acc (Array sh sh) -> Acc (Array sh a) -> Acc (Array sh b)
forall sh a b c.
(Shape sh, Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
zipWith Exp sh -> Exp a -> Exp b
f (Exp sh -> (Exp sh -> Exp sh) -> Acc (Array sh sh)
forall sh a.
(Shape sh, Elt a) =>
Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a)
generate (Acc (Array sh a) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh a)
xs) Exp sh -> Exp sh
forall a. a -> a
id) Acc (Array sh a)
xs

-- | Used to define the zipWith functions on more than two arrays
--
zipWithInduction
    :: (Shape sh, Elt a, Elt b)
    => ((Exp (a,b) -> rest) -> Acc (Array sh (a,b)) -> result) -- The zipWith function operating on one fewer array
    -> (Exp a -> Exp b -> rest)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> result
zipWithInduction :: ((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result
prev Exp a -> Exp b -> rest
f Acc (Array sh a)
as Acc (Array sh b)
bs = (Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result
prev (\(T2 Exp a
a Exp b
b) -> Exp a -> Exp b -> rest
f Exp a
a Exp b
b) (Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip Acc (Array sh a)
as Acc (Array sh b)
bs)


-- | Zip three arrays with the given function, analogous to 'zipWith'.
--
zipWith3
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d)
    => (Exp a -> Exp b -> Exp c -> Exp d)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
zipWith3 :: (Exp a -> Exp b -> Exp c -> Exp d)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
zipWith3 = ((Exp (a, b) -> Exp c -> Exp d)
 -> Acc (Array sh (a, b)) -> Acc (Array sh c) -> Acc (Array sh d))
-> (Exp a -> Exp b -> Exp c -> Exp d)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b) -> Exp c -> Exp d)
-> Acc (Array sh (a, b)) -> Acc (Array sh c) -> Acc (Array sh d)
forall sh a b c.
(Shape sh, Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
zipWith

-- | Zip four arrays with the given function, analogous to 'zipWith'.
--
zipWith4
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e)
    => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
zipWith4 :: (Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
zipWith4 = ((Exp (a, b) -> Exp c -> Exp d -> Exp e)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e))
-> (Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b) -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
forall sh a b c d.
(Shape sh, Elt a, Elt b, Elt c, Elt d) =>
(Exp a -> Exp b -> Exp c -> Exp d)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
zipWith3

-- | Zip five arrays with the given function, analogous to 'zipWith'.
--
zipWith5
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f)
    => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
zipWith5 :: (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
zipWith5 = ((Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f))
-> (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
forall sh a b c d e.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) =>
(Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
zipWith4

-- | Zip six arrays with the given function, analogous to 'zipWith'.
--
zipWith6
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g)
    => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
zipWith6 :: (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
zipWith6 = ((Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g))
-> (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
forall sh a b c d e f.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) =>
(Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
zipWith5

-- | Zip seven arrays with the given function, analogous to 'zipWith'.
--
zipWith7
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h)
    => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
    -> Acc (Array sh h)
zipWith7 :: (Exp a
 -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
zipWith7 = ((Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g)
 -> Acc (Array sh h))
-> (Exp a
    -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
forall sh a b c d e f g.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) =>
(Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
zipWith6

-- | Zip eight arrays with the given function, analogous to 'zipWith'.
--
zipWith8
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i)
    => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
    -> Acc (Array sh h)
    -> Acc (Array sh i)
zipWith8 :: (Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
zipWith8 = ((Exp (a, b)
  -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g)
 -> Acc (Array sh h)
 -> Acc (Array sh i))
-> (Exp a
    -> Exp b
    -> Exp c
    -> Exp d
    -> Exp e
    -> Exp f
    -> Exp g
    -> Exp h
    -> Exp i)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b)
 -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
forall sh a b c d e f g h.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g,
 Elt h) =>
(Exp a
 -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
zipWith7

-- | Zip nine arrays with the given function, analogous to 'zipWith'.
--
zipWith9
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i, Elt j)
    => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i -> Exp j)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
    -> Acc (Array sh h)
    -> Acc (Array sh i)
    -> Acc (Array sh j)
zipWith9 :: (Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i
 -> Exp j)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh j)
zipWith9 = ((Exp (a, b)
  -> Exp c
  -> Exp d
  -> Exp e
  -> Exp f
  -> Exp g
  -> Exp h
  -> Exp i
  -> Exp j)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g)
 -> Acc (Array sh h)
 -> Acc (Array sh i)
 -> Acc (Array sh j))
-> (Exp a
    -> Exp b
    -> Exp c
    -> Exp d
    -> Exp e
    -> Exp f
    -> Exp g
    -> Exp h
    -> Exp i
    -> Exp j)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh j)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
zipWithInduction (Exp (a, b)
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i
 -> Exp j)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh j)
forall sh a b c d e f g h i.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h,
 Elt i) =>
(Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
zipWith8


-- | Used to define the izipWith functions on two or more arrays
--
izipWithInduction
    :: (Shape sh, Elt a, Elt b)
    => ((Exp sh -> Exp (a,b) -> rest) -> Acc (Array sh (a,b)) -> result) -- The zipWith function operating on one fewer array
    -> (Exp sh -> Exp a -> Exp b -> rest)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> result
izipWithInduction :: ((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result
prev Exp sh -> Exp a -> Exp b -> rest
f Acc (Array sh a)
as Acc (Array sh b)
bs = (Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result
prev (\Exp sh
ix (T2 Exp a
a Exp b
b) -> Exp sh -> Exp a -> Exp b -> rest
f Exp sh
ix Exp a
a Exp b
b) (Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip Acc (Array sh a)
as Acc (Array sh b)
bs)


-- | Zip two arrays with a function that also takes the element index
--
izipWith
    :: (Shape sh, Elt a, Elt b, Elt c)
    => (Exp sh -> Exp a -> Exp b -> Exp c)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
izipWith :: (Exp sh -> Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
izipWith = ((Exp sh -> Exp (a, b) -> Exp c)
 -> Acc (Array sh (a, b)) -> Acc (Array sh c))
-> (Exp sh -> Exp a -> Exp b -> Exp c)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh -> Exp (a, b) -> Exp c)
-> Acc (Array sh (a, b)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp sh -> Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
imap

-- | Zip three arrays with a function that also takes the element index,
-- analogous to 'izipWith'.
--
izipWith3
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d)
    => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
izipWith3 :: (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
izipWith3 = ((Exp sh -> Exp (a, b) -> Exp c -> Exp d)
 -> Acc (Array sh (a, b)) -> Acc (Array sh c) -> Acc (Array sh d))
-> (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh -> Exp (a, b) -> Exp c -> Exp d)
-> Acc (Array sh (a, b)) -> Acc (Array sh c) -> Acc (Array sh d)
forall sh a b c.
(Shape sh, Elt a, Elt b, Elt c) =>
(Exp sh -> Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
izipWith

-- | Zip four arrays with the given function that also takes the element index,
-- analogous to 'zipWith'.
--
izipWith4
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e)
    => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
izipWith4 :: (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
izipWith4 = ((Exp sh -> Exp (a, b) -> Exp c -> Exp d -> Exp e)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e))
-> (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh -> Exp (a, b) -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
forall sh a b c d.
(Shape sh, Elt a, Elt b, Elt c, Elt d) =>
(Exp sh -> Exp a -> Exp b -> Exp c -> Exp d)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
izipWith3

-- | Zip five arrays with the given function that also takes the element index,
-- analogous to 'zipWith'.
--
izipWith5
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f)
    => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
izipWith5 :: (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
izipWith5 = ((Exp sh -> Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f))
-> (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh -> Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
forall sh a b c d e.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) =>
(Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
izipWith4

-- | Zip six arrays with the given function that also takes the element index,
-- analogous to 'zipWith'.
--
izipWith6
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g)
    => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
izipWith6 :: (Exp sh
 -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
izipWith6 = ((Exp sh
  -> Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g))
-> (Exp sh
    -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh -> Exp (a, b) -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
forall sh a b c d e f.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) =>
(Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
izipWith5

-- | Zip seven arrays with the given function that also takes the element
-- index, analogous to 'zipWith'.
--
izipWith7
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h)
    => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
    -> Acc (Array sh h)
izipWith7 :: (Exp sh
 -> Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
izipWith7 = ((Exp sh
  -> Exp (a, b)
  -> Exp c
  -> Exp d
  -> Exp e
  -> Exp f
  -> Exp g
  -> Exp h)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g)
 -> Acc (Array sh h))
-> (Exp sh
    -> Exp a
    -> Exp b
    -> Exp c
    -> Exp d
    -> Exp e
    -> Exp f
    -> Exp g
    -> Exp h)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh
 -> Exp (a, b)
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
forall sh a b c d e f g.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) =>
(Exp sh
 -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
izipWith6

-- | Zip eight arrays with the given function that also takes the element
-- index, analogous to 'zipWith'.
--
izipWith8
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i)
    => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
    -> Acc (Array sh h)
    -> Acc (Array sh i)
izipWith8 :: (Exp sh
 -> Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
izipWith8 = ((Exp sh
  -> Exp (a, b)
  -> Exp c
  -> Exp d
  -> Exp e
  -> Exp f
  -> Exp g
  -> Exp h
  -> Exp i)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g)
 -> Acc (Array sh h)
 -> Acc (Array sh i))
-> (Exp sh
    -> Exp a
    -> Exp b
    -> Exp c
    -> Exp d
    -> Exp e
    -> Exp f
    -> Exp g
    -> Exp h
    -> Exp i)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh
 -> Exp (a, b)
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
forall sh a b c d e f g h.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g,
 Elt h) =>
(Exp sh
 -> Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
izipWith7

-- | Zip nine arrays with the given function that also takes the element index,
-- analogous to 'zipWith'.
--
izipWith9
    :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i, Elt j)
    => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i -> Exp j)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh c)
    -> Acc (Array sh d)
    -> Acc (Array sh e)
    -> Acc (Array sh f)
    -> Acc (Array sh g)
    -> Acc (Array sh h)
    -> Acc (Array sh i)
    -> Acc (Array sh j)
izipWith9 :: (Exp sh
 -> Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i
 -> Exp j)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh j)
izipWith9 = ((Exp sh
  -> Exp (a, b)
  -> Exp c
  -> Exp d
  -> Exp e
  -> Exp f
  -> Exp g
  -> Exp h
  -> Exp i
  -> Exp j)
 -> Acc (Array sh (a, b))
 -> Acc (Array sh c)
 -> Acc (Array sh d)
 -> Acc (Array sh e)
 -> Acc (Array sh f)
 -> Acc (Array sh g)
 -> Acc (Array sh h)
 -> Acc (Array sh i)
 -> Acc (Array sh j))
-> (Exp sh
    -> Exp a
    -> Exp b
    -> Exp c
    -> Exp d
    -> Exp e
    -> Exp f
    -> Exp g
    -> Exp h
    -> Exp i
    -> Exp j)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh j)
forall sh a b rest result.
(Shape sh, Elt a, Elt b) =>
((Exp sh -> Exp (a, b) -> rest) -> Acc (Array sh (a, b)) -> result)
-> (Exp sh -> Exp a -> Exp b -> rest)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> result
izipWithInduction (Exp sh
 -> Exp (a, b)
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i
 -> Exp j)
-> Acc (Array sh (a, b))
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh j)
forall sh a b c d e f g h i.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h,
 Elt i) =>
(Exp sh
 -> Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
izipWith8


-- | Combine the elements of two arrays pairwise. The shape of the result is the
-- intersection of the two argument shapes.
--
-- >>> let m1 = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> let m2 = fromList (Z:.10:.5) [0..] :: Matrix Float
-- >>> run $ zip (use m1) (use m2)
-- Matrix (Z :. 5 :. 5)
--   [   (0,0.0),   (1,1.0),   (2,2.0),   (3,3.0),   (4,4.0),
--      (10,5.0),  (11,6.0),  (12,7.0),  (13,8.0),  (14,9.0),
--     (20,10.0), (21,11.0), (22,12.0), (23,13.0), (24,14.0),
--     (30,15.0), (31,16.0), (32,17.0), (33,18.0), (34,19.0),
--     (40,20.0), (41,21.0), (42,22.0), (43,23.0), (44,24.0)]
--
zip :: (Shape sh, Elt a, Elt b)
    => Acc (Array sh a)
    -> Acc (Array sh b)
    -> Acc (Array sh (a, b))
zip :: Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip = (Exp a -> Exp b -> Exp (a, b))
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
forall sh a b c.
(Shape sh, Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
zipWith Exp a -> Exp b -> Exp (a, b)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2

-- | Take three arrays and return an array of triples, analogous to zip.
--
zip3 :: (Shape sh, Elt a, Elt b, Elt c)
     => Acc (Array sh a)
     -> Acc (Array sh b)
     -> Acc (Array sh c)
     -> Acc (Array sh (a, b, c))
zip3 :: Acc (Array sh a)
-> Acc (Array sh b) -> Acc (Array sh c) -> Acc (Array sh (a, b, c))
zip3 = (Exp a -> Exp b -> Exp c -> Exp (a, b, c))
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh (a, b, c))
forall sh a b c d.
(Shape sh, Elt a, Elt b, Elt c, Elt d) =>
(Exp a -> Exp b -> Exp c -> Exp d)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
zipWith3 Exp a -> Exp b -> Exp c -> Exp (a, b, c)
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con x0 -> con x1 -> con x2 -> con (x0, x1, x2)
T3

-- | Take four arrays and return an array of quadruples, analogous to zip.
--
zip4 :: (Shape sh, Elt a, Elt b, Elt c, Elt d)
     => Acc (Array sh a)
     -> Acc (Array sh b)
     -> Acc (Array sh c)
     -> Acc (Array sh d)
     -> Acc (Array sh (a, b, c, d))
zip4 :: Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh (a, b, c, d))
zip4 = (Exp a -> Exp b -> Exp c -> Exp d -> Exp (a, b, c, d))
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh (a, b, c, d))
forall sh a b c d e.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) =>
(Exp a -> Exp b -> Exp c -> Exp d -> Exp e)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
zipWith4 Exp a -> Exp b -> Exp c -> Exp d -> Exp (a, b, c, d)
forall (con :: * -> *) x0 x1 x2 x3.
IsPattern con (x0, x1, x2, x3) (con x0, con x1, con x2, con x3) =>
con x0 -> con x1 -> con x2 -> con x3 -> con (x0, x1, x2, x3)
T4

-- | Take five arrays and return an array of five-tuples, analogous to zip.
--
zip5 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e)
     => Acc (Array sh a)
     -> Acc (Array sh b)
     -> Acc (Array sh c)
     -> Acc (Array sh d)
     -> Acc (Array sh e)
     -> Acc (Array sh (a, b, c, d, e))
zip5 :: Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh (a, b, c, d, e))
zip5 = (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp (a, b, c, d, e))
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh (a, b, c, d, e))
forall sh a b c d e f.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) =>
(Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
zipWith5 Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp (a, b, c, d, e)
forall (con :: * -> *) x0 x1 x2 x3 x4.
IsPattern
  con
  (x0, x1, x2, x3, x4)
  (con x0, con x1, con x2, con x3, con x4) =>
con x0
-> con x1 -> con x2 -> con x3 -> con x4 -> con (x0, x1, x2, x3, x4)
T5

-- | Take six arrays and return an array of six-tuples, analogous to zip.
--
zip6 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f)
     => Acc (Array sh a)
     -> Acc (Array sh b)
     -> Acc (Array sh c)
     -> Acc (Array sh d)
     -> Acc (Array sh e)
     -> Acc (Array sh f)
     -> Acc (Array sh (a, b, c, d, e, f))
zip6 :: Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh (a, b, c, d, e, f))
zip6 = (Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp (a, b, c, d, e, f))
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh (a, b, c, d, e, f))
forall sh a b c d e f g.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) =>
(Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
zipWith6 Exp a
-> Exp b
-> Exp c
-> Exp d
-> Exp e
-> Exp f
-> Exp (a, b, c, d, e, f)
forall (con :: * -> *) x0 x1 x2 x3 x4 x5.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5)
  (con x0, con x1, con x2, con x3, con x4, con x5) =>
con x0
-> con x1
-> con x2
-> con x3
-> con x4
-> con x5
-> con (x0, x1, x2, x3, x4, x5)
T6

-- | Take seven arrays and return an array of seven-tuples, analogous to zip.
--
zip7 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g)
     => Acc (Array sh a)
     -> Acc (Array sh b)
     -> Acc (Array sh c)
     -> Acc (Array sh d)
     -> Acc (Array sh e)
     -> Acc (Array sh f)
     -> Acc (Array sh g)
     -> Acc (Array sh (a, b, c, d, e, f, g))
zip7 :: Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh (a, b, c, d, e, f, g))
zip7 = (Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp (a, b, c, d, e, f, g))
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh (a, b, c, d, e, f, g))
forall sh a b c d e f g h.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g,
 Elt h) =>
(Exp a
 -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
zipWith7 Exp a
-> Exp b
-> Exp c
-> Exp d
-> Exp e
-> Exp f
-> Exp g
-> Exp (a, b, c, d, e, f, g)
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con x0
-> con x1
-> con x2
-> con x3
-> con x4
-> con x5
-> con x6
-> con (x0, x1, x2, x3, x4, x5, x6)
T7

-- | Take seven arrays and return an array of seven-tuples, analogous to zip.
--
zip8 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h)
     => Acc (Array sh a)
     -> Acc (Array sh b)
     -> Acc (Array sh c)
     -> Acc (Array sh d)
     -> Acc (Array sh e)
     -> Acc (Array sh f)
     -> Acc (Array sh g)
     -> Acc (Array sh h)
     -> Acc (Array sh (a, b, c, d, e, f, g, h))
zip8 :: Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh (a, b, c, d, e, f, g, h))
zip8 = (Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp (a, b, c, d, e, f, g, h))
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh (a, b, c, d, e, f, g, h))
forall sh a b c d e f g h i.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h,
 Elt i) =>
(Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
zipWith8 Exp a
-> Exp b
-> Exp c
-> Exp d
-> Exp e
-> Exp f
-> Exp g
-> Exp h
-> Exp (a, b, c, d, e, f, g, h)
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con x0
-> con x1
-> con x2
-> con x3
-> con x4
-> con x5
-> con x6
-> con x7
-> con (x0, x1, x2, x3, x4, x5, x6, x7)
T8

-- | Take seven arrays and return an array of seven-tuples, analogous to zip.
--
zip9 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i)
     => Acc (Array sh a)
     -> Acc (Array sh b)
     -> Acc (Array sh c)
     -> Acc (Array sh d)
     -> Acc (Array sh e)
     -> Acc (Array sh f)
     -> Acc (Array sh g)
     -> Acc (Array sh h)
     -> Acc (Array sh i)
     -> Acc (Array sh (a, b, c, d, e, f, g, h, i))
zip9 :: Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i))
zip9 = (Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i
 -> Exp (a, b, c, d, e, f, g, h, i))
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i))
forall sh a b c d e f g h i j.
(Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h,
 Elt i, Elt j) =>
(Exp a
 -> Exp b
 -> Exp c
 -> Exp d
 -> Exp e
 -> Exp f
 -> Exp g
 -> Exp h
 -> Exp i
 -> Exp j)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
-> Acc (Array sh d)
-> Acc (Array sh e)
-> Acc (Array sh f)
-> Acc (Array sh g)
-> Acc (Array sh h)
-> Acc (Array sh i)
-> Acc (Array sh j)
zipWith9 Exp a
-> Exp b
-> Exp c
-> Exp d
-> Exp e
-> Exp f
-> Exp g
-> Exp h
-> Exp i
-> Exp (a, b, c, d, e, f, g, h, i)
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con x0
-> con x1
-> con x2
-> con x3
-> con x4
-> con x5
-> con x6
-> con x7
-> con x8
-> con (x0, x1, x2, x3, x4, x5, x6, x7, x8)
T9


-- | The converse of 'zip', but the shape of the two results is identical to the
-- shape of the argument.
--
-- If the argument array is manifest in memory, 'unzip' is a no-op.
--
unzip :: (Shape sh, Elt a, Elt b)
      => Acc (Array sh (a, b))
      -> (Acc (Array sh a), Acc (Array sh b))
unzip :: Acc (Array sh (a, b)) -> (Acc (Array sh a), Acc (Array sh b))
unzip Acc (Array sh (a, b))
arr = ((Exp (a, b) -> Exp a) -> Acc (Array sh (a, b)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b) -> Exp a
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp a
fst Acc (Array sh (a, b))
arr, (Exp (a, b) -> Exp b) -> Acc (Array sh (a, b)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b) -> Exp b
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp b
snd Acc (Array sh (a, b))
arr)

-- | Take an array of triples and return three arrays, analogous to 'unzip'.
--
unzip3 :: (Shape sh, Elt a, Elt b, Elt c)
       => Acc (Array sh (a, b, c))
       -> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c))
unzip3 :: Acc (Array sh (a, b, c))
-> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c))
unzip3 Acc (Array sh (a, b, c))
xs = ((Exp (a, b, c) -> Exp a)
-> Acc (Array sh (a, b, c)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c) -> Exp a
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con (x0, x1, x2) -> con x0
get1 Acc (Array sh (a, b, c))
xs, (Exp (a, b, c) -> Exp b)
-> Acc (Array sh (a, b, c)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c) -> Exp b
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con (x0, x1, x2) -> con x1
get2 Acc (Array sh (a, b, c))
xs, (Exp (a, b, c) -> Exp c)
-> Acc (Array sh (a, b, c)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c) -> Exp c
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con (x0, x1, x2) -> con x2
get3 Acc (Array sh (a, b, c))
xs)
  where
    get1 :: con (x0, x1, x2) -> con x0
get1 (T3 con x0
a con x1
_ con x2
_) = con x0
a
    get2 :: con (x0, x1, x2) -> con x1
get2 (T3 con x0
_ con x1
b con x2
_) = con x1
b
    get3 :: con (x0, x1, x2) -> con x2
get3 (T3 con x0
_ con x1
_ con x2
c) = con x2
c


-- | Take an array of quadruples and return four arrays, analogous to 'unzip'.
--
unzip4 :: (Shape sh, Elt a, Elt b, Elt c, Elt d)
       => Acc (Array sh (a, b, c, d))
       -> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c), Acc (Array sh d))
unzip4 :: Acc (Array sh (a, b, c, d))
-> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c),
    Acc (Array sh d))
unzip4 Acc (Array sh (a, b, c, d))
xs = ((Exp (a, b, c, d) -> Exp a)
-> Acc (Array sh (a, b, c, d)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d) -> Exp a
forall (con :: * -> *) x0 x1 x2 x3.
IsPattern con (x0, x1, x2, x3) (con x0, con x1, con x2, con x3) =>
con (x0, x1, x2, x3) -> con x0
get1 Acc (Array sh (a, b, c, d))
xs, (Exp (a, b, c, d) -> Exp b)
-> Acc (Array sh (a, b, c, d)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d) -> Exp b
forall (con :: * -> *) x0 x1 x2 x3.
IsPattern con (x0, x1, x2, x3) (con x0, con x1, con x2, con x3) =>
con (x0, x1, x2, x3) -> con x1
get2 Acc (Array sh (a, b, c, d))
xs, (Exp (a, b, c, d) -> Exp c)
-> Acc (Array sh (a, b, c, d)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d) -> Exp c
forall (con :: * -> *) x0 x1 x2 x3.
IsPattern con (x0, x1, x2, x3) (con x0, con x1, con x2, con x3) =>
con (x0, x1, x2, x3) -> con x2
get3 Acc (Array sh (a, b, c, d))
xs, (Exp (a, b, c, d) -> Exp d)
-> Acc (Array sh (a, b, c, d)) -> Acc (Array sh d)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d) -> Exp d
forall (con :: * -> *) x0 x1 x2 x3.
IsPattern con (x0, x1, x2, x3) (con x0, con x1, con x2, con x3) =>
con (x0, x1, x2, x3) -> con x3
get4 Acc (Array sh (a, b, c, d))
xs)
  where
    get1 :: con (x0, x1, x2, x3) -> con x0
get1 (T4 con x0
a con x1
_ con x2
_ con x3
_) = con x0
a
    get2 :: con (x0, x1, x2, x3) -> con x1
get2 (T4 con x0
_ con x1
b con x2
_ con x3
_) = con x1
b
    get3 :: con (x0, x1, x2, x3) -> con x2
get3 (T4 con x0
_ con x1
_ con x2
c con x3
_) = con x2
c
    get4 :: con (x0, x1, x2, x3) -> con x3
get4 (T4 con x0
_ con x1
_ con x2
_ con x3
d) = con x3
d

-- | Take an array of 5-tuples and return five arrays, analogous to 'unzip'.
--
unzip5 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e)
       => Acc (Array sh (a, b, c, d, e))
       -> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c), Acc (Array sh d), Acc (Array sh e))
unzip5 :: Acc (Array sh (a, b, c, d, e))
-> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c),
    Acc (Array sh d), Acc (Array sh e))
unzip5 Acc (Array sh (a, b, c, d, e))
xs = ((Exp (a, b, c, d, e) -> Exp a)
-> Acc (Array sh (a, b, c, d, e)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e) -> Exp a
forall (con :: * -> *) x0 x1 x2 x3 x4.
IsPattern
  con
  (x0, x1, x2, x3, x4)
  (con x0, con x1, con x2, con x3, con x4) =>
con (x0, x1, x2, x3, x4) -> con x0
get1 Acc (Array sh (a, b, c, d, e))
xs, (Exp (a, b, c, d, e) -> Exp b)
-> Acc (Array sh (a, b, c, d, e)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e) -> Exp b
forall (con :: * -> *) x0 x1 x2 x3 x4.
IsPattern
  con
  (x0, x1, x2, x3, x4)
  (con x0, con x1, con x2, con x3, con x4) =>
con (x0, x1, x2, x3, x4) -> con x1
get2 Acc (Array sh (a, b, c, d, e))
xs, (Exp (a, b, c, d, e) -> Exp c)
-> Acc (Array sh (a, b, c, d, e)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e) -> Exp c
forall (con :: * -> *) x0 x1 x2 x3 x4.
IsPattern
  con
  (x0, x1, x2, x3, x4)
  (con x0, con x1, con x2, con x3, con x4) =>
con (x0, x1, x2, x3, x4) -> con x2
get3 Acc (Array sh (a, b, c, d, e))
xs, (Exp (a, b, c, d, e) -> Exp d)
-> Acc (Array sh (a, b, c, d, e)) -> Acc (Array sh d)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e) -> Exp d
forall (con :: * -> *) x0 x1 x2 x3 x4.
IsPattern
  con
  (x0, x1, x2, x3, x4)
  (con x0, con x1, con x2, con x3, con x4) =>
con (x0, x1, x2, x3, x4) -> con x3
get4 Acc (Array sh (a, b, c, d, e))
xs, (Exp (a, b, c, d, e) -> Exp e)
-> Acc (Array sh (a, b, c, d, e)) -> Acc (Array sh e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e) -> Exp e
forall (con :: * -> *) x0 x1 x2 x3 x4.
IsPattern
  con
  (x0, x1, x2, x3, x4)
  (con x0, con x1, con x2, con x3, con x4) =>
con (x0, x1, x2, x3, x4) -> con x4
get5 Acc (Array sh (a, b, c, d, e))
xs)
  where
    get1 :: con (x0, x1, x2, x3, x4) -> con x0
get1 (T5 con x0
a con x1
_ con x2
_ con x3
_ con x4
_) = con x0
a
    get2 :: con (x0, x1, x2, x3, x4) -> con x1
get2 (T5 con x0
_ con x1
b con x2
_ con x3
_ con x4
_) = con x1
b
    get3 :: con (x0, x1, x2, x3, x4) -> con x2
get3 (T5 con x0
_ con x1
_ con x2
c con x3
_ con x4
_) = con x2
c
    get4 :: con (x0, x1, x2, x3, x4) -> con x3
get4 (T5 con x0
_ con x1
_ con x2
_ con x3
d con x4
_) = con x3
d
    get5 :: con (x0, x1, x2, x3, x4) -> con x4
get5 (T5 con x0
_ con x1
_ con x2
_ con x3
_ con x4
e) = con x4
e

-- | Take an array of 6-tuples and return six arrays, analogous to 'unzip'.
--
unzip6 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f)
       => Acc (Array sh (a, b, c, d, e, f))
       -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c)
          , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f))
unzip6 :: Acc (Array sh (a, b, c, d, e, f))
-> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c),
    Acc (Array sh d), Acc (Array sh e), Acc (Array sh f))
unzip6 Acc (Array sh (a, b, c, d, e, f))
xs = ((Exp (a, b, c, d, e, f) -> Exp a)
-> Acc (Array sh (a, b, c, d, e, f)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f) -> Exp a
forall (con :: * -> *) x0 x1 x2 x3 x4 x5.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5)
  (con x0, con x1, con x2, con x3, con x4, con x5) =>
con (x0, x1, x2, x3, x4, x5) -> con x0
get1 Acc (Array sh (a, b, c, d, e, f))
xs, (Exp (a, b, c, d, e, f) -> Exp b)
-> Acc (Array sh (a, b, c, d, e, f)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f) -> Exp b
forall (con :: * -> *) x0 x1 x2 x3 x4 x5.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5)
  (con x0, con x1, con x2, con x3, con x4, con x5) =>
con (x0, x1, x2, x3, x4, x5) -> con x1
get2 Acc (Array sh (a, b, c, d, e, f))
xs, (Exp (a, b, c, d, e, f) -> Exp c)
-> Acc (Array sh (a, b, c, d, e, f)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f) -> Exp c
forall (con :: * -> *) x0 x1 x2 x3 x4 x5.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5)
  (con x0, con x1, con x2, con x3, con x4, con x5) =>
con (x0, x1, x2, x3, x4, x5) -> con x2
get3 Acc (Array sh (a, b, c, d, e, f))
xs, (Exp (a, b, c, d, e, f) -> Exp d)
-> Acc (Array sh (a, b, c, d, e, f)) -> Acc (Array sh d)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f) -> Exp d
forall (con :: * -> *) x0 x1 x2 x3 x4 x5.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5)
  (con x0, con x1, con x2, con x3, con x4, con x5) =>
con (x0, x1, x2, x3, x4, x5) -> con x3
get4 Acc (Array sh (a, b, c, d, e, f))
xs, (Exp (a, b, c, d, e, f) -> Exp e)
-> Acc (Array sh (a, b, c, d, e, f)) -> Acc (Array sh e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f) -> Exp e
forall (con :: * -> *) x0 x1 x2 x3 x4 x5.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5)
  (con x0, con x1, con x2, con x3, con x4, con x5) =>
con (x0, x1, x2, x3, x4, x5) -> con x4
get5 Acc (Array sh (a, b, c, d, e, f))
xs, (Exp (a, b, c, d, e, f) -> Exp f)
-> Acc (Array sh (a, b, c, d, e, f)) -> Acc (Array sh f)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f) -> Exp f
forall (con :: * -> *) x0 x1 x2 x3 x4 x5.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5)
  (con x0, con x1, con x2, con x3, con x4, con x5) =>
con (x0, x1, x2, x3, x4, x5) -> con x5
get6 Acc (Array sh (a, b, c, d, e, f))
xs)
  where
    get1 :: con (x0, x1, x2, x3, x4, x5) -> con x0
get1 (T6 con x0
a con x1
_ con x2
_ con x3
_ con x4
_ con x5
_) = con x0
a
    get2 :: con (x0, x1, x2, x3, x4, x5) -> con x1
get2 (T6 con x0
_ con x1
b con x2
_ con x3
_ con x4
_ con x5
_) = con x1
b
    get3 :: con (x0, x1, x2, x3, x4, x5) -> con x2
get3 (T6 con x0
_ con x1
_ con x2
c con x3
_ con x4
_ con x5
_) = con x2
c
    get4 :: con (x0, x1, x2, x3, x4, x5) -> con x3
get4 (T6 con x0
_ con x1
_ con x2
_ con x3
d con x4
_ con x5
_) = con x3
d
    get5 :: con (x0, x1, x2, x3, x4, x5) -> con x4
get5 (T6 con x0
_ con x1
_ con x2
_ con x3
_ con x4
e con x5
_) = con x4
e
    get6 :: con (x0, x1, x2, x3, x4, x5) -> con x5
get6 (T6 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
f) = con x5
f

-- | Take an array of 7-tuples and return seven arrays, analogous to 'unzip'.
--
unzip7 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g)
       => Acc (Array sh (a, b, c, d, e, f, g))
       -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c)
          , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f)
          , Acc (Array sh g))
unzip7 :: Acc (Array sh (a, b, c, d, e, f, g))
-> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c),
    Acc (Array sh d), Acc (Array sh e), Acc (Array sh f),
    Acc (Array sh g))
unzip7 Acc (Array sh (a, b, c, d, e, f, g))
xs = ( (Exp (a, b, c, d, e, f, g) -> Exp a)
-> Acc (Array sh (a, b, c, d, e, f, g)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g) -> Exp a
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con (x0, x1, x2, x3, x4, x5, x6) -> con x0
get1 Acc (Array sh (a, b, c, d, e, f, g))
xs, (Exp (a, b, c, d, e, f, g) -> Exp b)
-> Acc (Array sh (a, b, c, d, e, f, g)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g) -> Exp b
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con (x0, x1, x2, x3, x4, x5, x6) -> con x1
get2 Acc (Array sh (a, b, c, d, e, f, g))
xs, (Exp (a, b, c, d, e, f, g) -> Exp c)
-> Acc (Array sh (a, b, c, d, e, f, g)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g) -> Exp c
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con (x0, x1, x2, x3, x4, x5, x6) -> con x2
get3 Acc (Array sh (a, b, c, d, e, f, g))
xs
            , (Exp (a, b, c, d, e, f, g) -> Exp d)
-> Acc (Array sh (a, b, c, d, e, f, g)) -> Acc (Array sh d)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g) -> Exp d
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con (x0, x1, x2, x3, x4, x5, x6) -> con x3
get4 Acc (Array sh (a, b, c, d, e, f, g))
xs, (Exp (a, b, c, d, e, f, g) -> Exp e)
-> Acc (Array sh (a, b, c, d, e, f, g)) -> Acc (Array sh e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g) -> Exp e
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con (x0, x1, x2, x3, x4, x5, x6) -> con x4
get5 Acc (Array sh (a, b, c, d, e, f, g))
xs, (Exp (a, b, c, d, e, f, g) -> Exp f)
-> Acc (Array sh (a, b, c, d, e, f, g)) -> Acc (Array sh f)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g) -> Exp f
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con (x0, x1, x2, x3, x4, x5, x6) -> con x5
get6 Acc (Array sh (a, b, c, d, e, f, g))
xs
            , (Exp (a, b, c, d, e, f, g) -> Exp g)
-> Acc (Array sh (a, b, c, d, e, f, g)) -> Acc (Array sh g)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g) -> Exp g
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6) =>
con (x0, x1, x2, x3, x4, x5, x6) -> con x6
get7 Acc (Array sh (a, b, c, d, e, f, g))
xs )
  where
    get1 :: con (x0, x1, x2, x3, x4, x5, x6) -> con x0
get1 (T7 con x0
a con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
_) = con x0
a
    get2 :: con (x0, x1, x2, x3, x4, x5, x6) -> con x1
get2 (T7 con x0
_ con x1
b con x2
_ con x3
_ con x4
_ con x5
_ con x6
_) = con x1
b
    get3 :: con (x0, x1, x2, x3, x4, x5, x6) -> con x2
get3 (T7 con x0
_ con x1
_ con x2
c con x3
_ con x4
_ con x5
_ con x6
_) = con x2
c
    get4 :: con (x0, x1, x2, x3, x4, x5, x6) -> con x3
get4 (T7 con x0
_ con x1
_ con x2
_ con x3
d con x4
_ con x5
_ con x6
_) = con x3
d
    get5 :: con (x0, x1, x2, x3, x4, x5, x6) -> con x4
get5 (T7 con x0
_ con x1
_ con x2
_ con x3
_ con x4
e con x5
_ con x6
_) = con x4
e
    get6 :: con (x0, x1, x2, x3, x4, x5, x6) -> con x5
get6 (T7 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
f con x6
_) = con x5
f
    get7 :: con (x0, x1, x2, x3, x4, x5, x6) -> con x6
get7 (T7 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
g) = con x6
g

-- | Take an array of 8-tuples and return eight arrays, analogous to 'unzip'.
--
unzip8 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h)
       => Acc (Array sh (a, b, c, d, e, f, g, h))
       -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c)
          , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f)
          , Acc (Array sh g), Acc (Array sh h) )
unzip8 :: Acc (Array sh (a, b, c, d, e, f, g, h))
-> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c),
    Acc (Array sh d), Acc (Array sh e), Acc (Array sh f),
    Acc (Array sh g), Acc (Array sh h))
unzip8 Acc (Array sh (a, b, c, d, e, f, g, h))
xs = ( (Exp (a, b, c, d, e, f, g, h) -> Exp a)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp a
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x0
get1 Acc (Array sh (a, b, c, d, e, f, g, h))
xs, (Exp (a, b, c, d, e, f, g, h) -> Exp b)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp b
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x1
get2 Acc (Array sh (a, b, c, d, e, f, g, h))
xs, (Exp (a, b, c, d, e, f, g, h) -> Exp c)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp c
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x2
get3 Acc (Array sh (a, b, c, d, e, f, g, h))
xs
            , (Exp (a, b, c, d, e, f, g, h) -> Exp d)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh d)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp d
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x3
get4 Acc (Array sh (a, b, c, d, e, f, g, h))
xs, (Exp (a, b, c, d, e, f, g, h) -> Exp e)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp e
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x4
get5 Acc (Array sh (a, b, c, d, e, f, g, h))
xs, (Exp (a, b, c, d, e, f, g, h) -> Exp f)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh f)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp f
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x5
get6 Acc (Array sh (a, b, c, d, e, f, g, h))
xs
            , (Exp (a, b, c, d, e, f, g, h) -> Exp g)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh g)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp g
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x6
get7 Acc (Array sh (a, b, c, d, e, f, g, h))
xs, (Exp (a, b, c, d, e, f, g, h) -> Exp h)
-> Acc (Array sh (a, b, c, d, e, f, g, h)) -> Acc (Array sh h)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h) -> Exp h
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7) =>
con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x7
get8 Acc (Array sh (a, b, c, d, e, f, g, h))
xs )
  where
    get1 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x0
get1 (T8 con x0
a con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
_ con x7
_) = con x0
a
    get2 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x1
get2 (T8 con x0
_ con x1
b con x2
_ con x3
_ con x4
_ con x5
_ con x6
_ con x7
_) = con x1
b
    get3 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x2
get3 (T8 con x0
_ con x1
_ con x2
c con x3
_ con x4
_ con x5
_ con x6
_ con x7
_) = con x2
c
    get4 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x3
get4 (T8 con x0
_ con x1
_ con x2
_ con x3
d con x4
_ con x5
_ con x6
_ con x7
_) = con x3
d
    get5 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x4
get5 (T8 con x0
_ con x1
_ con x2
_ con x3
_ con x4
e con x5
_ con x6
_ con x7
_) = con x4
e
    get6 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x5
get6 (T8 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
f con x6
_ con x7
_) = con x5
f
    get7 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x6
get7 (T8 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
g con x7
_) = con x6
g
    get8 :: con (x0, x1, x2, x3, x4, x5, x6, x7) -> con x7
get8 (T8 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
_ con x7
h) = con x7
h

-- | Take an array of 9-tuples and return nine arrays, analogous to 'unzip'.
--
unzip9 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i)
       => Acc (Array sh (a, b, c, d, e, f, g, h, i))
       -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c)
          , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f)
          , Acc (Array sh g), Acc (Array sh h), Acc (Array sh i))
unzip9 :: Acc (Array sh (a, b, c, d, e, f, g, h, i))
-> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c),
    Acc (Array sh d), Acc (Array sh e), Acc (Array sh f),
    Acc (Array sh g), Acc (Array sh h), Acc (Array sh i))
unzip9 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs = ( (Exp (a, b, c, d, e, f, g, h, i) -> Exp a)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp a
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x0
get1 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs, (Exp (a, b, c, d, e, f, g, h, i) -> Exp b)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh b)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp b
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x1
get2 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs, (Exp (a, b, c, d, e, f, g, h, i) -> Exp c)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh c)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp c
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x2
get3 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs
            , (Exp (a, b, c, d, e, f, g, h, i) -> Exp d)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh d)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp d
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x3
get4 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs, (Exp (a, b, c, d, e, f, g, h, i) -> Exp e)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp e
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x4
get5 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs, (Exp (a, b, c, d, e, f, g, h, i) -> Exp f)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh f)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp f
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x5
get6 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs
            , (Exp (a, b, c, d, e, f, g, h, i) -> Exp g)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh g)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp g
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x6
get7 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs, (Exp (a, b, c, d, e, f, g, h, i) -> Exp h)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh h)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp h
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x7
get8 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs, (Exp (a, b, c, d, e, f, g, h, i) -> Exp i)
-> Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> Acc (Array sh i)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (a, b, c, d, e, f, g, h, i) -> Exp i
forall (con :: * -> *) x0 x1 x2 x3 x4 x5 x6 x7 x8.
IsPattern
  con
  (x0, x1, x2, x3, x4, x5, x6, x7, x8)
  (con x0, con x1, con x2, con x3, con x4, con x5, con x6, con x7,
   con x8) =>
con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x8
get9 Acc (Array sh (a, b, c, d, e, f, g, h, i))
xs )
  where
    get1 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x0
get1 (T9 con x0
a con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
_ con x7
_ con x8
_) = con x0
a
    get2 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x1
get2 (T9 con x0
_ con x1
b con x2
_ con x3
_ con x4
_ con x5
_ con x6
_ con x7
_ con x8
_) = con x1
b
    get3 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x2
get3 (T9 con x0
_ con x1
_ con x2
c con x3
_ con x4
_ con x5
_ con x6
_ con x7
_ con x8
_) = con x2
c
    get4 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x3
get4 (T9 con x0
_ con x1
_ con x2
_ con x3
d con x4
_ con x5
_ con x6
_ con x7
_ con x8
_) = con x3
d
    get5 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x4
get5 (T9 con x0
_ con x1
_ con x2
_ con x3
_ con x4
e con x5
_ con x6
_ con x7
_ con x8
_) = con x4
e
    get6 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x5
get6 (T9 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
f con x6
_ con x7
_ con x8
_) = con x5
f
    get7 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x6
get7 (T9 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
g con x7
_ con x8
_) = con x6
g
    get8 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x7
get8 (T9 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
_ con x7
h con x8
_) = con x7
h
    get9 :: con (x0, x1, x2, x3, x4, x5, x6, x7, x8) -> con x8
get9 (T9 con x0
_ con x1
_ con x2
_ con x3
_ con x4
_ con x5
_ con x6
_ con x7
_ con x8
i) = con x8
i


-- Reductions
-- ----------

-- | Reduction of an array of arbitrary rank to a single scalar value. The first
-- argument needs to be an /associative/ function to enable efficient parallel
-- implementation. The initial element does not need to be an identity element.
--
-- >>> let vec = fromList (Z:.10) [0..] :: Vector Float
-- >>> run $ foldAll (+) 42 (use vec)
-- Scalar Z [87.0]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Float
-- >>> run $ foldAll (+) 0 (use mat)
-- Scalar Z [1225.0]
--
foldAll
    :: (Shape sh, Elt a)
    => (Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array sh a)
    -> Acc (Scalar a)
foldAll :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array sh a) -> Acc (Scalar a)
foldAll Exp a -> Exp a -> Exp a
f Exp a
e Acc (Array sh a)
arr = (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (DIM0 :. Int) a) -> Acc (Scalar a)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold Exp a -> Exp a -> Exp a
f Exp a
e (Acc (Array sh a) -> Acc (Array (DIM0 :. Int) a)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Acc (Vector e)
flatten Acc (Array sh a)
arr)

-- | Variant of 'foldAll' that requires the reduced array to be non-empty and
-- does not need a default value. The first argument must be an /associative/
-- function.
--
fold1All
    :: (Shape sh, Elt a)
    => (Exp a -> Exp a -> Exp a)
    -> Acc (Array sh a)
    -> Acc (Scalar a)
fold1All :: (Exp a -> Exp a -> Exp a) -> Acc (Array sh a) -> Acc (Scalar a)
fold1All Exp a -> Exp a -> Exp a
f Acc (Array sh a)
arr = (Exp a -> Exp a -> Exp a)
-> Acc (Array (DIM0 :. Int) a) -> Acc (Scalar a)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold1 Exp a -> Exp a -> Exp a
f (Acc (Array sh a) -> Acc (Array (DIM0 :. Int) a)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Acc (Vector e)
flatten Acc (Array sh a)
arr)


-- | Segmented reduction along the innermost dimension of an array. The segment
-- descriptor specifies the lengths of the logical sub-arrays, each of which is
-- reduced independently. The innermost dimension must contain at least as many
-- elements as required by the segment descriptor (sum thereof).
--
-- >>> let seg = fromList (Z:.4) [1,4,0,3] :: Segments Int
-- >>> seg
-- Vector (Z :. 4) [1,4,0,3]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ foldSeg (+) 0 (use mat) (use seg)
-- Matrix (Z :. 5 :. 4)
--   [  0,  10, 0,  18,
--     10,  50, 0,  48,
--     20,  90, 0,  78,
--     30, 130, 0, 108,
--     40, 170, 0, 138]
--
foldSeg
    :: forall sh e i. (Shape sh, Elt e, Elt i, i ~ EltR i, IsIntegral i)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
foldSeg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
foldSeg Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg = (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh a i.
(Shape sh, Elt a, Elt i, IsIntegral i, i ~ EltR i) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) a)
foldSeg' Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr ((Exp i -> Exp i -> Exp i)
-> Exp i -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl Exp i -> Exp i -> Exp i
plus Exp i
zero Acc (Segments i)
seg)
  where
    (Exp i -> Exp i -> Exp i
plus, Exp i
zero) =
      case IsIntegral i => IntegralType i
forall a. IsIntegral a => IntegralType a
integralType @i of
        TypeInt{}    -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt8{}   -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt16{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt32{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt64{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord{}   -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord8{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord16{} -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord32{} -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord64{} -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)


-- | Variant of 'foldSeg' that requires /all/ segments of the reduced array
-- to be non-empty, and does not need a default value. The segment
-- descriptor species the length of each of the logical sub-arrays.
--
fold1Seg
    :: forall sh e i. (Shape sh, Elt e, Elt i, i ~ EltR i, IsIntegral i)
    => (Exp e -> Exp e -> Exp e)
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
fold1Seg :: (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
fold1Seg Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg = (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh a i.
(Shape sh, Elt a, Elt i, IsIntegral i, i ~ EltR i) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) a)
fold1Seg' Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
arr ((Exp i -> Exp i -> Exp i)
-> Exp i -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl Exp i -> Exp i -> Exp i
plus Exp i
zero Acc (Segments i)
seg)
  where
    plus :: Exp i -> Exp i -> Exp i
    zero :: Exp i
    (Exp i -> Exp i -> Exp i
plus, Exp i
zero) =
      case IsIntegral (EltR i) => IntegralType (EltR i)
forall a. IsIntegral a => IntegralType a
integralType @(EltR i) of
        TypeInt{}    -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt8{}   -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt16{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt32{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeInt64{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord{}   -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord8{}  -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord16{} -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord32{} -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)
        TypeWord64{} -> (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+), Exp i
0)


-- Specialised reductions
-- ----------------------
--
-- Leave the results of these as scalar arrays to make it clear that these are
-- array computations, and thus can not be nested.

-- | Check if all elements along the innermost dimension satisfy a predicate.
--
-- >>> let mat = fromList (Z :. 4 :. 10) [1,2,3,4,5,6,7,8,9,10,1,1,1,1,1,2,2,2,2,2,2,4,6,8,10,12,14,16,18,20,1,3,5,7,9,11,13,15,17,19] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 4 :. 10)
--   [ 1, 2, 3, 4,  5,  6,  7,  8,  9, 10,
--     1, 1, 1, 1,  1,  2,  2,  2,  2,  2,
--     2, 4, 6, 8, 10, 12, 14, 16, 18, 20,
--     1, 3, 5, 7,  9, 11, 13, 15, 17, 19]
--
-- >>> run $ all even (use mat)
-- Vector (Z :. 4) [False,False,True,False]
--
all :: (Shape sh, Elt e)
    => (Exp e -> Exp Bool)
    -> Acc (Array (sh:.Int) e)
    -> Acc (Array sh Bool)
all :: (Exp e -> Exp Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Array sh Bool)
all Exp e -> Exp Bool
f = Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
forall sh.
Shape sh =>
Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
and (Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool))
-> (Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) Bool))
-> Acc (Array (sh :. Int) e)
-> Acc (Array sh Bool)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp e -> Exp Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) Bool)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp e -> Exp Bool
f

-- | Check if any element along the innermost dimension satisfies the predicate.
--
-- >>> let mat = fromList (Z :. 4 :. 10) [1,2,3,4,5,6,7,8,9,10,1,1,1,1,1,2,2,2,2,2,2,4,6,8,10,12,14,16,18,20,1,3,5,7,9,11,13,15,17,19] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 4 :. 10)
--   [ 1, 2, 3, 4,  5,  6,  7,  8,  9, 10,
--     1, 1, 1, 1,  1,  2,  2,  2,  2,  2,
--     2, 4, 6, 8, 10, 12, 14, 16, 18, 20,
--     1, 3, 5, 7,  9, 11, 13, 15, 17, 19]
--
-- >>> run $ any even (use mat)
-- Vector (Z :. 4) [True,True,True,False]
--
any :: (Shape sh, Elt e)
    => (Exp e -> Exp Bool)
    -> Acc (Array (sh:.Int) e)
    -> Acc (Array sh Bool)
any :: (Exp e -> Exp Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Array sh Bool)
any Exp e -> Exp Bool
f = Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
forall sh.
Shape sh =>
Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
or (Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool))
-> (Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) Bool))
-> Acc (Array (sh :. Int) e)
-> Acc (Array sh Bool)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp e -> Exp Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) Bool)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp e -> Exp Bool
f

-- | Check if all elements along the innermost dimension are 'True'.
--
and :: Shape sh
    => Acc (Array (sh:.Int) Bool)
    -> Acc (Array sh Bool)
and :: Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
and = (Exp Bool -> Exp Bool -> Exp Bool)
-> Exp Bool -> Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold Exp Bool -> Exp Bool -> Exp Bool
(&&) Exp Bool
HasCallStack => Exp Bool
True_

-- | Check if any element along the innermost dimension is 'True'.
--
or :: Shape sh
   => Acc (Array (sh:.Int) Bool)
   -> Acc (Array sh Bool)
or :: Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
or = (Exp Bool -> Exp Bool -> Exp Bool)
-> Exp Bool -> Acc (Array (sh :. Int) Bool) -> Acc (Array sh Bool)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold Exp Bool -> Exp Bool -> Exp Bool
(||) Exp Bool
HasCallStack => Exp Bool
False_

-- | Compute the sum of elements along the innermost dimension of the array. To
-- find the sum of the entire array, 'flatten' it first.
--
-- >>> let mat = fromList (Z:.2:.5) [0..] :: Matrix Int
-- >>> run $ sum (use mat)
-- Vector (Z :. 2) [10,35]
--
sum :: (Shape sh, Num e)
    => Acc (Array (sh:.Int) e)
    -> Acc (Array sh e)
sum :: Acc (Array (sh :. Int) e) -> Acc (Array sh e)
sum = (Exp e -> Exp e -> Exp e)
-> Exp e -> Acc (Array (sh :. Int) e) -> Acc (Array sh e)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
(+) Exp e
0

-- | Compute the product of the elements along the innermost dimension of the
-- array. To find the product of the entire array, 'flatten' it first.
--
-- >>> let mat = fromList (Z:.2:.5) [0..] :: Matrix Int
-- >>> run $ product (use mat)
-- Vector (Z :. 2) [0,15120]
--
product
    :: (Shape sh, Num e)
    => Acc (Array (sh:.Int) e)
    -> Acc (Array sh e)
product :: Acc (Array (sh :. Int) e) -> Acc (Array sh e)
product = (Exp e -> Exp e -> Exp e)
-> Exp e -> Acc (Array (sh :. Int) e) -> Acc (Array sh e)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
(*) Exp e
1

-- | Yield the minimum element along the innermost dimension of the array. To
-- find find the minimum element of the entire array, 'flatten' it first.
--
-- The array must not be empty. See also 'fold1'.
--
-- >>> let mat = fromList (Z :. 3 :. 4) [1,4,3,8, 0,2,8,4, 7,9,8,8] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 3 :. 4)
--   [ 1, 4, 3, 8,
--     0, 2, 8, 4,
--     7, 9, 8, 8]
--
-- >>> run $ minimum (use mat)
-- Vector (Z :. 3) [1,0,7]
--
minimum
    :: (Shape sh, Ord e)
    => Acc (Array (sh:.Int) e)
    -> Acc (Array sh e)
minimum :: Acc (Array (sh :. Int) e) -> Acc (Array sh e)
minimum = (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e) -> Acc (Array sh e)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold1 Exp e -> Exp e -> Exp e
forall a. Ord a => Exp a -> Exp a -> Exp a
min

-- | Yield the maximum element along the innermost dimension of the array. To
-- find the maximum element of the entire array, 'flatten' it first.
--
-- The array must not be empty. See also 'fold1'.
--
-- >>> let mat = fromList (Z :. 3 :. 4) [1,4,3,8, 0,2,8,4, 7,9,8,8] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 3 :. 4)
--   [ 1, 4, 3, 8,
--     0, 2, 8, 4,
--     7, 9, 8, 8]
--
-- >>> run $ maximum (use mat)
-- Vector (Z :. 3) [8,8,9]
--
maximum
    :: (Shape sh, Ord e)
    => Acc (Array (sh:.Int) e)
    -> Acc (Array sh e)
maximum :: Acc (Array (sh :. Int) e) -> Acc (Array sh e)
maximum = (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e) -> Acc (Array sh e)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold1 Exp e -> Exp e -> Exp e
forall a. Ord a => Exp a -> Exp a -> Exp a
max


-- Composite scans
-- ---------------

-- | Left-to-right pre-scan (aka exclusive scan). As for 'scan', the first
-- argument must be an /associative/ function. Denotationally, we have:
--
-- > prescanl f e = afst . scanl' f e
--
-- >>> let vec = fromList (Z:.10) [1..10] :: Vector Int
-- >>> run $ prescanl (+) 0 (use vec)
-- Vector (Z :. 10) [0,1,3,6,10,15,21,28,36,45]
--
prescanl
    :: (Shape sh, Elt a)
    => (Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array (sh:.Int) a)
    -> Acc (Array (sh:.Int) a)
prescanl :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
prescanl Exp a -> Exp a -> Exp a
f Exp a
e = Acc (Array (sh :. Int) a, Array sh a) -> Acc (Array (sh :. Int) a)
forall a b. (Arrays a, Arrays b) => Acc (a, b) -> Acc a
afst (Acc (Array (sh :. Int) a, Array sh a)
 -> Acc (Array (sh :. Int) a))
-> (Acc (Array (sh :. Int) a)
    -> Acc (Array (sh :. Int) a, Array sh a))
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanl' Exp a -> Exp a -> Exp a
f Exp a
e

-- | Left-to-right post-scan, a variant of 'scanl1' with an initial value. As
-- with 'scanl1', the array must not be empty. Denotationally, we have:
--
-- > postscanl f e = map (e `f`) . scanl1 f
--
-- >>> let vec = fromList (Z:.10) [1..10] :: Vector Int
-- >>> run $ postscanl (+) 42 (use vec)
-- Vector (Z :. 10) [43,45,48,52,57,63,70,78,87,97]
--
postscanl
    :: (Shape sh, Elt a)
    => (Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array (sh:.Int) a)
    -> Acc (Array (sh:.Int) a)
postscanl :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
postscanl Exp a -> Exp a -> Exp a
f Exp a
e = (Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp a
e Exp a -> Exp a -> Exp a
`f`) (Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a))
-> (Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 Exp a -> Exp a -> Exp a
f

-- | Right-to-left pre-scan (aka exclusive scan). As for 'scan', the first
-- argument must be an /associative/ function. Denotationally, we have:
--
-- > prescanr f e = afst . scanr' f e
--
prescanr
    :: (Shape sh, Elt a)
    => (Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array (sh:.Int) a)
    -> Acc (Array (sh:.Int) a)
prescanr :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
prescanr Exp a -> Exp a -> Exp a
f Exp a
e = Acc (Array (sh :. Int) a, Array sh a) -> Acc (Array (sh :. Int) a)
forall a b. (Arrays a, Arrays b) => Acc (a, b) -> Acc a
afst (Acc (Array (sh :. Int) a, Array sh a)
 -> Acc (Array (sh :. Int) a))
-> (Acc (Array (sh :. Int) a)
    -> Acc (Array (sh :. Int) a, Array sh a))
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanr' Exp a -> Exp a -> Exp a
f Exp a
e

-- | Right-to-left postscan, a variant of 'scanr1' with an initial value.
-- Denotationally, we have:
--
-- > postscanr f e = map (e `f`) . scanr1 f
--
postscanr
    :: (Shape sh, Elt a)
    => (Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array (sh:.Int) a)
    -> Acc (Array (sh:.Int) a)
postscanr :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
postscanr Exp a -> Exp a -> Exp a
f Exp a
e = (Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp a -> Exp a -> Exp a
`f` Exp a
e) (Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a))
-> (Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanr1 Exp a -> Exp a -> Exp a
f


-- Segmented scans
-- ---------------

-- | Segmented version of 'scanl' along the innermost dimension of an array. The
-- innermost dimension must have at least as many elements as the sum of the
-- segment descriptor.
--
-- >>> let seg = fromList (Z:.4) [1,4,0,3] :: Segments Int
-- >>> seg
-- Vector (Z :. 4) [1,4,0,3]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ scanlSeg (+) 0 (use mat) (use seg)
-- Matrix (Z :. 5 :. 12)
--   [ 0,  0, 0,  1,  3,   6,  10, 0, 0,  5, 11,  18,
--     0, 10, 0, 11, 23,  36,  50, 0, 0, 15, 31,  48,
--     0, 20, 0, 21, 43,  66,  90, 0, 0, 25, 51,  78,
--     0, 30, 0, 31, 63,  96, 130, 0, 0, 35, 71, 108,
--     0, 40, 0, 41, 83, 126, 170, 0, 0, 45, 91, 138]
--
scanlSeg
    :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
scanlSeg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanlSeg Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg =
  if Acc (Array (sh :. Int) e) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Array (sh :. Int) e)
arr Exp Bool -> Exp Bool -> Exp Bool
|| Acc (Segments i) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Segments i)
flags
    then Exp (sh :. Int) -> Exp e -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp sh
sh Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
sz Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg) Exp e
z
    else (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanl1Seg Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
arr' Acc (Segments i)
seg'
  where
    -- Segmented exclusive scan is implemented by first injecting the seed
    -- element at the head of each segment, and then performing a segmented
    -- inclusive scan.
    --
    -- This is done by creating a vector entirely of the seed element, and
    -- overlaying the input data in all places other than at the start of
    -- a segment.
    --
    Exp sh
sh ::. Exp Int
sz = Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr
    seg' :: Acc (Segments i)
seg'      = (Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
+Exp i
1) Acc (Segments i)
seg
    arr' :: Acc (Array (sh :. Int) e)
arr'      = (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> (Exp (sh :. Int) -> Exp (Maybe (sh :. Int)))
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp e -> Exp e -> Exp e
forall a b. a -> b -> a
const
                        (Exp (sh :. Int) -> Exp e -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp sh
sh Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
sz Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg) Exp e
z)
                        (\(Exp sh
sx ::. Exp Int
i) -> Exp (sh :. Int) -> Exp (Maybe (sh :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp sh
sx Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
i Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral (Acc (Segments i)
inc Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
i)))
                        (Exp Int -> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp Int -> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
take (Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
flags) Acc (Array (sh :. Int) e)
arr)

    -- Each element in the segments must be shifted to the right one additional
    -- place for each successive segment, to make room for the seed element.
    -- Here, we make use of the fact that the vector returned by 'mkHeadFlags'
    -- contains non-unit entries, which indicate zero length segments.
    --
    flags :: Acc (Segments i)
flags     = Acc (Segments i) -> Acc (Segments i)
forall i.
(Integral i, FromIntegral i Int) =>
Acc (Segments i) -> Acc (Segments i)
mkHeadFlags Acc (Segments i)
seg
    inc :: Acc (Segments i)
inc       = (Exp i -> Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Acc (Segments i)
flags


-- | Segmented version of 'scanl'' along the innermost dimension of an array. The
-- innermost dimension must have at least as many elements as the sum of the
-- segment descriptor.
--
-- The first element of the resulting tuple is a vector of scanned values. The
-- second element is a vector of segment scan totals and has the same size as
-- the segment vector.
--
-- >>> let seg = fromList (Z:.4) [1,4,0,3] :: Segments Int
-- >>> seg
-- Vector (Z :. 4) [1,4,0,3]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> let (res,sums) = run $ scanl'Seg (+) 0 (use mat) (use seg)
-- >>> res
-- Matrix (Z :. 5 :. 8)
--   [ 0, 0,  1,  3,   6, 0,  5, 11,
--     0, 0, 11, 23,  36, 0, 15, 31,
--     0, 0, 21, 43,  66, 0, 25, 51,
--     0, 0, 31, 63,  96, 0, 35, 71,
--     0, 0, 41, 83, 126, 0, 45, 91]
-- >>> sums
-- Matrix (Z :. 5 :. 4)
--   [  0,  10, 0,  18,
--     10,  50, 0,  48,
--     20,  90, 0,  78,
--     30, 130, 0, 108,
--     40, 170, 0, 138]
--
scanl'Seg
    :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e, Array (sh:.Int) e)
scanl'Seg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
scanl'Seg Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg =
  if Acc (Array (sh :. Int) e) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Array (sh :. Int) e)
arr
    then Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Array (sh :. Int) e)
arr  (Exp (sh :. Int) -> Exp e -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr) Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg) Exp e
z)
    else Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Array (sh :. Int) e)
body Acc (Array (sh :. Int) e)
sums
  where
    -- Segmented scan' is implemented by deconstructing a segmented exclusive
    -- scan, to separate the final value and scan body.
    --
    -- TLM: Segmented scans, and this version in particular, expend a lot of
    --      effort scanning flag arrays. On inspection it appears that several
    --      of these operations are duplicated, but this will not be picked up
    --      by sharing _observation_. Perhaps a global CSE-style pass would be
    --      beneficial.
    --
    arr' :: Acc (Array (sh :. Int) e)
arr'        = (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanlSeg Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg

    -- Extract the final reduction value for each segment, which is at the last
    -- index of each segment.
    --
    seg' :: Acc (Segments i)
seg'        = (Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
+Exp i
1) Acc (Segments i)
seg
    tails :: Acc (Segments i)
tails       = (Exp i -> Exp i -> Exp i)
-> Acc (Segments i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a b c.
(Shape sh, Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
zipWith Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Acc (Segments i)
seg (Acc (Segments i) -> Acc (Segments i))
-> Acc (Segments i) -> Acc (Segments i)
forall a b. (a -> b) -> a -> b
$ (Exp i -> Exp i -> Exp i)
-> Exp i -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
prescanl Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Exp i
0 Acc (Segments i)
seg'
    sums :: Acc (Array (sh :. Int) e)
sums        = Exp (sh :. Int)
-> (Exp (sh :. Int) -> Exp (sh :. Int))
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute
                    (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr') Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg)
                    (\(Exp sh
sz ::. Exp Int
i) -> Exp sh
sz Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral (Acc (Segments i)
tails Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
i))
                    Acc (Array (sh :. Int) e)
arr'

    -- Slice out the body of each segment.
    --
    -- Build a head-flags representation based on the original segment
    -- descriptor. This contains the target length of each of the body segments,
    -- which is one fewer element than the actual bodies stored in arr'. Thus,
    -- the flags align with the last element of each body section, and when
    -- scanned, this element will be incremented over.
    --
    offset :: Acc (Segments i)
offset      = (Exp i -> Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Acc (Segments i)
seg
    inc :: Acc (Segments i)
inc         = (Exp i -> Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+)
                (Acc (Segments i) -> Acc (Segments i))
-> Acc (Segments i) -> Acc (Segments i)
forall a b. (a -> b) -> a -> b
$ (Exp i -> Exp i -> Exp i)
-> Acc (Segments i)
-> (Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Segments i)
-> Acc (Segments i)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) (Exp (DIM0 :. Int) -> Exp i -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Exp Int -> Exp (DIM0 :. Int)) -> Exp Int -> Exp (DIM0 :. Int)
forall a b. (a -> b) -> a -> b
$ Acc (Array (sh :. Int) e) -> Exp Int
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
size Acc (Array (sh :. Int) e)
arr Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Exp Int
1) Exp i
0)
                              (\Exp (DIM0 :. Int)
ix -> Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp i -> Exp (DIM0 :. Int)
forall i.
(Integral i, FromIntegral i Int) =>
Exp i -> Exp (DIM0 :. Int)
index1' (Acc (Segments i)
offset Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp (DIM0 :. Int)
ix)))
                              (Exp (DIM0 :. Int) -> Exp i -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Acc (Segments i) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Segments i)
seg) (Exp i
1 :: Exp i))

    len :: Exp i
len         = Acc (Segments i)
offset Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
offset Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
- Exp Int
1)
    body :: Acc (Array (sh :. Int) e)
body        = Exp (sh :. Int)
-> (Exp (sh :. Int) -> Exp (sh :. Int))
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute
                    (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr) Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral Exp i
len)
                    (\(Exp sh
sz ::. Exp Int
i) -> Exp sh
sz Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
i Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral (Acc (Segments i)
inc Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
i))
                    Acc (Array (sh :. Int) e)
arr'


-- | Segmented version of 'scanl1' along the innermost dimension.
--
-- As with 'scanl1', the total number of elements considered, in this case given
-- by the 'sum' of segment descriptor, must not be zero. The input vector must
-- contain at least this many elements.
--
-- Zero length segments are allowed, and the behaviour is as if those entries
-- were not present in the segment descriptor; that is:
--
-- > scanl1Seg f xs [n,0,0] == scanl1Seg f xs [n]   where n /= 0
--
-- >>> let seg = fromList (Z:.4) [1,4,0,3] :: Segments Int
-- >>> seg
-- Vector (Z :. 4) [1,4,0,3]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ scanl1Seg (+) (use mat) (use seg)
-- Matrix (Z :. 5 :. 8)
--   [  0,  1,  3,   6,  10,  5, 11,  18,
--     10, 11, 23,  36,  50, 15, 31,  48,
--     20, 21, 43,  66,  90, 25, 51,  78,
--     30, 31, 63,  96, 130, 35, 71, 108,
--     40, 41, 83, 126, 170, 45, 91, 138]
--
scanl1Seg
    :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
scanl1Seg :: (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanl1Seg Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg
  = (Exp (i, e) -> Exp e)
-> Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (i, e) -> Exp e
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp b
snd
  (Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e))
-> (Acc (Array (sh :. Int) (i, e))
    -> Acc (Array (sh :. Int) (i, e)))
-> Acc (Array (sh :. Int) (i, e))
-> Acc (Array (sh :. Int) e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp (i, e) -> Exp (i, e) -> Exp (i, e))
-> Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) (i, e))
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 ((Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
forall e i.
(Elt e, Num i, Bits i) =>
(Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
segmentedL Exp e -> Exp e -> Exp e
f)
  (Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e))
-> Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e)
forall a b. (a -> b) -> a -> b
$ Acc (Array (sh :. Int) i)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) (i, e))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip (Exp (sh :. All)
-> Acc (Array (SliceShape (sh :. All)) i)
-> Acc (Array (FullShape (sh :. All)) i)
forall slix e.
(Slice slix, Elt e) =>
Exp slix
-> Acc (Array (SliceShape slix) e)
-> Acc (Array (FullShape slix) e)
replicate ((Exp sh :. All) -> Exp (Plain (Exp sh :. All))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr) Exp sh -> All -> Exp sh :. All
forall tail head. tail -> head -> tail :. head
:. All
All)) (Acc (Segments i) -> Acc (Segments i)
forall i.
(Integral i, FromIntegral i Int) =>
Acc (Segments i) -> Acc (Segments i)
mkHeadFlags Acc (Segments i)
seg)) Acc (Array (sh :. Int) e)
arr

-- |Segmented version of 'prescanl'.
--
prescanlSeg
    :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
prescanlSeg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
prescanlSeg Exp e -> Exp e -> Exp e
f Exp e
e Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg
  = Acc (Array (sh :. Int) e, Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall a b. (Arrays a, Arrays b) => Acc (a, b) -> Acc a
afst
  (Acc (Array (sh :. Int) e, Array (sh :. Int) e)
 -> Acc (Array (sh :. Int) e))
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall a b. (a -> b) -> a -> b
$ (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
scanl'Seg Exp e -> Exp e -> Exp e
f Exp e
e Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg

-- |Segmented version of 'postscanl'.
--
postscanlSeg
    :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
postscanlSeg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
postscanlSeg Exp e -> Exp e -> Exp e
f Exp e
e Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg
  = (Exp e -> Exp e)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp e -> Exp e -> Exp e
f Exp e
e)
  (Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e))
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall a b. (a -> b) -> a -> b
$ (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanl1Seg Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg

-- | Segmented version of 'scanr' along the innermost dimension of an array. The
-- innermost dimension must have at least as many elements as the sum of the
-- segment descriptor.
--
-- >>> let seg = fromList (Z:.4) [1,4,0,3] :: Segments Int
-- >>> seg
-- Vector (Z :. 4) [1,4,0,3]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ scanrSeg (+) 0 (use mat) (use seg)
-- Matrix (Z :. 5 :. 12)
--   [  2, 0,  18,  15, 11,  6, 0, 0,  24, 17,  9, 0,
--     12, 0,  58,  45, 31, 16, 0, 0,  54, 37, 19, 0,
--     22, 0,  98,  75, 51, 26, 0, 0,  84, 57, 29, 0,
--     32, 0, 138, 105, 71, 36, 0, 0, 114, 77, 39, 0,
--     42, 0, 178, 135, 91, 46, 0, 0, 144, 97, 49, 0]
--
scanrSeg
    :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
scanrSeg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanrSeg Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg =
  if Acc (Array (sh :. Int) e) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Array (sh :. Int) e)
arr Exp Bool -> Exp Bool -> Exp Bool
|| Acc (Segments i) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Segments i)
flags
    then Exp (sh :. Int) -> Exp e -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp sh
sh Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
sz Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg) Exp e
z
    else (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanr1Seg Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
arr' Acc (Segments i)
seg'
  where
    Exp sh
sh ::. Exp Int
sz    = Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr

    -- Using technique described for 'scanlSeg', where we intersperse the array
    -- with the seed element at the start of each segment, and then perform an
    -- inclusive segmented scan.
    --
    flags :: Acc (Segments i)
flags       = Acc (Segments i) -> Acc (Segments i)
forall i.
(Integral i, FromIntegral i Int) =>
Acc (Segments i) -> Acc (Segments i)
mkHeadFlags Acc (Segments i)
seg
    inc :: Acc (Segments i)
inc         = (Exp i -> Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Acc (Segments i)
flags

    seg' :: Acc (Segments i)
seg'        = (Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
+Exp i
1) Acc (Segments i)
seg
    arr' :: Acc (Array (sh :. Int) e)
arr'        = (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> (Exp (sh :. Int) -> Exp (Maybe (sh :. Int)))
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp e -> Exp e -> Exp e
forall a b. a -> b -> a
const
                          (Exp (sh :. Int) -> Exp e -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp sh
sh Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
sz Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg) Exp e
z)
                          (\(Exp sh
sx ::. Exp Int
i) -> Exp (sh :. Int) -> Exp (Maybe (sh :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp sh
sx Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
i Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral (Acc (Segments i)
inc Acc (Segments i) -> Exp Int -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp Int -> Exp e
!! Exp Int
i) Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
- Exp Int
1))
                          (Exp Int -> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp Int -> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
drop (Exp Int
sz Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
- Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
flags) Acc (Array (sh :. Int) e)
arr)


-- | Segmented version of 'scanr''.
--
-- >>> let seg = fromList (Z:.4) [1,4,0,3] :: Segments Int
-- >>> seg
-- Vector (Z :. 4) [1,4,0,3]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> let (res,sums) = run $ scanr'Seg (+) 0 (use mat) (use seg)
-- >>> res
-- Matrix (Z :. 5 :. 8)
--   [ 0,  15, 11,  6, 0, 17,  9, 0,
--     0,  45, 31, 16, 0, 37, 19, 0,
--     0,  75, 51, 26, 0, 57, 29, 0,
--     0, 105, 71, 36, 0, 77, 39, 0,
--     0, 135, 91, 46, 0, 97, 49, 0]
-- >>> sums
-- Matrix (Z :. 5 :. 4)
--   [  2,  18, 0,  24,
--     12,  58, 0,  54,
--     22,  98, 0,  84,
--     32, 138, 0, 114,
--     42, 178, 0, 144]
--
scanr'Seg
    :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e, Array (sh:.Int) e)
scanr'Seg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
scanr'Seg Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg =
  if Acc (Array (sh :. Int) e) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Array (sh :. Int) e)
arr
    then Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Array (sh :. Int) e)
arr  (Exp (sh :. Int) -> Exp e -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr) Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg) Exp e
z)
    else Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Array (sh :. Int) e)
body Acc (Array (sh :. Int) e)
sums
  where
    -- Using technique described for scanl'Seg
    --
    arr' :: Acc (Array (sh :. Int) e)
arr'        = (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanrSeg Exp e -> Exp e -> Exp e
f Exp e
z Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg

    -- reduction values
    seg' :: Acc (Segments i)
seg'        = (Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
+Exp i
1) Acc (Segments i)
seg
    heads :: Acc (Segments i)
heads       = (Exp i -> Exp i -> Exp i)
-> Exp i -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
prescanl Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Exp i
0 Acc (Segments i)
seg'
    sums :: Acc (Array (sh :. Int) e)
sums        = Exp (sh :. Int)
-> (Exp (sh :. Int) -> Exp (sh :. Int))
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute
                    (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr') Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Acc (Segments i) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Segments i)
seg)
                    (\(Exp sh
sz ::.Exp Int
i) -> Exp sh
sz Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral (Acc (Segments i)
heads Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
i))
                    Acc (Array (sh :. Int) e)
arr'

    -- body segments
    flags :: Acc (Segments i)
flags       = Acc (Segments i) -> Acc (Segments i)
forall i.
(Integral i, FromIntegral i Int) =>
Acc (Segments i) -> Acc (Segments i)
mkHeadFlags Acc (Segments i)
seg
    inc :: Acc (Segments i)
inc         = (Exp i -> Exp i -> Exp i) -> Acc (Segments i) -> Acc (Segments i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Acc (Segments i)
flags
    body :: Acc (Array (sh :. Int) e)
body        = Exp (sh :. Int)
-> (Exp (sh :. Int) -> Exp (sh :. Int))
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute
                    (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr) Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp (DIM0 :. Int) -> Exp Int
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp a
indexHead (Acc (Segments i) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Segments i)
flags))
                    (\(Exp sh
sz ::. Exp Int
i) -> Exp sh
sz Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
i Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral (Acc (Segments i)
inc Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
i))
                    Acc (Array (sh :. Int) e)
arr'


-- | Segmented version of 'scanr1'.
--
-- >>> let seg = fromList (Z:.4) [1,4,0,3] :: Segments Int
-- >>> seg
-- Vector (Z :. 4) [1,4,0,3]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ scanr1Seg (+) (use mat) (use seg)
-- Matrix (Z :. 5 :. 8)
--   [  0,  10,   9,  7,  4,  18, 13,  7,
--     10,  50,  39, 27, 14,  48, 33, 17,
--     20,  90,  69, 47, 24,  78, 53, 27,
--     30, 130,  99, 67, 34, 108, 73, 37,
--     40, 170, 129, 87, 44, 138, 93, 47]
--
scanr1Seg
    :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
scanr1Seg :: (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanr1Seg Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
arr Acc (Segments i)
seg
  = (Exp (i, e) -> Exp e)
-> Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (i, e) -> Exp e
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp b
snd
  (Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e))
-> (Acc (Array (sh :. Int) (i, e))
    -> Acc (Array (sh :. Int) (i, e)))
-> Acc (Array (sh :. Int) (i, e))
-> Acc (Array (sh :. Int) e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp (i, e) -> Exp (i, e) -> Exp (i, e))
-> Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) (i, e))
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanr1 ((Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
forall e i.
(Elt e, Num i, Bits i) =>
(Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
segmentedR Exp e -> Exp e -> Exp e
f)
  (Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e))
-> Acc (Array (sh :. Int) (i, e)) -> Acc (Array (sh :. Int) e)
forall a b. (a -> b) -> a -> b
$ Acc (Array (sh :. Int) i)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) (i, e))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip (Exp (sh :. All)
-> Acc (Array (SliceShape (sh :. All)) i)
-> Acc (Array (FullShape (sh :. All)) i)
forall slix e.
(Slice slix, Elt e) =>
Exp slix
-> Acc (Array (SliceShape slix) e)
-> Acc (Array (FullShape slix) e)
replicate ((Exp sh :. All) -> Exp (Plain (Exp sh :. All))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr) Exp sh -> All -> Exp sh :. All
forall tail head. tail -> head -> tail :. head
:. All
All)) (Acc (Segments i) -> Acc (Segments i)
forall i.
(Integral i, FromIntegral i Int) =>
Acc (Segments i) -> Acc (Segments i)
mkTailFlags Acc (Segments i)
seg)) Acc (Array (sh :. Int) e)
arr


-- |Segmented version of 'prescanr'.
--
prescanrSeg
    :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
prescanrSeg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
prescanrSeg Exp e -> Exp e -> Exp e
f Exp e
e Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg
  = Acc (Array (sh :. Int) e, Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall a b. (Arrays a, Arrays b) => Acc (a, b) -> Acc a
afst
  (Acc (Array (sh :. Int) e, Array (sh :. Int) e)
 -> Acc (Array (sh :. Int) e))
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall a b. (a -> b) -> a -> b
$ (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e, Array (sh :. Int) e)
scanr'Seg Exp e -> Exp e -> Exp e
f Exp e
e Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg

-- |Segmented version of 'postscanr'.
--
postscanrSeg
    :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int)
    => (Exp e -> Exp e -> Exp e)
    -> Exp e
    -> Acc (Array (sh:.Int) e)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) e)
postscanrSeg :: (Exp e -> Exp e -> Exp e)
-> Exp e
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
postscanrSeg Exp e -> Exp e -> Exp e
f Exp e
e Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg
  = (Exp e -> Exp e)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp e -> Exp e -> Exp e
f Exp e
e)
  (Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e))
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall a b. (a -> b) -> a -> b
$ (Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
forall sh e i.
(Shape sh, Slice sh, Elt e, Integral i, Bits i,
 FromIntegral i Int) =>
(Exp e -> Exp e -> Exp e)
-> Acc (Array (sh :. Int) e)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) e)
scanr1Seg Exp e -> Exp e -> Exp e
f Acc (Array (sh :. Int) e)
vec Acc (Segments i)
seg


-- Segmented scan helpers
-- ----------------------

-- | Compute head flags vector from segment vector for left-scans.
--
-- The vector will be full of zeros in the body of a segment, and non-zero
-- otherwise. The "flag" value, if greater than one, indicates that several
-- empty segments are represented by this single flag entry. This is additional
-- data is used by exclusive segmented scan.
--
mkHeadFlags
    :: (Integral i, FromIntegral i Int)
    => Acc (Segments i)
    -> Acc (Segments i)
mkHeadFlags :: Acc (Segments i) -> Acc (Segments i)
mkHeadFlags Acc (Segments i)
seg
  = Acc (Segments i) -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
init
  (Acc (Segments i) -> Acc (Segments i))
-> Acc (Segments i) -> Acc (Segments i)
forall a b. (a -> b) -> a -> b
$ (Exp i -> Exp i -> Exp i)
-> Acc (Segments i)
-> (Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Segments i)
-> Acc (Segments i)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Acc (Segments i)
zeros (\Exp (DIM0 :. Int)
ix -> Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp i -> Exp (DIM0 :. Int)
forall i.
(Integral i, FromIntegral i Int) =>
Exp i -> Exp (DIM0 :. Int)
index1' (Acc (Segments i)
offset Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp (DIM0 :. Int)
ix))) Acc (Segments i)
ones
  where
    T2 Acc (Segments i)
offset Acc (Array DIM0 i)
len = (Exp i -> Exp i -> Exp i)
-> Exp i -> Acc (Segments i) -> Acc (Segments i, Array DIM0 i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanl' Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Exp i
0 Acc (Segments i)
seg
    zeros :: Acc (Segments i)
zeros         = Exp (DIM0 :. Int) -> Exp i -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp i -> Exp (DIM0 :. Int)
forall i.
(Integral i, FromIntegral i Int) =>
Exp i -> Exp (DIM0 :. Int)
index1' (Exp i -> Exp (DIM0 :. Int)) -> Exp i -> Exp (DIM0 :. Int)
forall a b. (a -> b) -> a -> b
$ Acc (Array DIM0 i) -> Exp i
forall e. Elt e => Acc (Scalar e) -> Exp e
the Acc (Array DIM0 i)
len Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
+ Exp i
1) Exp i
0
    ones :: Acc (Segments i)
ones          = Exp (DIM0 :. Int) -> Exp i -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
index1  (Exp Int -> Exp (DIM0 :. Int)) -> Exp Int -> Exp (DIM0 :. Int)
forall a b. (a -> b) -> a -> b
$ Acc (Segments i) -> Exp Int
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
size Acc (Segments i)
offset) Exp i
1

-- | Compute tail flags vector from segment vector for right-scans. That
-- is, the flag is placed at the last place in each segment.
--
mkTailFlags
    :: (Integral i, FromIntegral i Int)
    => Acc (Segments i)
    -> Acc (Segments i)
mkTailFlags :: Acc (Segments i) -> Acc (Segments i)
mkTailFlags Acc (Segments i)
seg
  = Acc (Segments i) -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
init
  (Acc (Segments i) -> Acc (Segments i))
-> Acc (Segments i) -> Acc (Segments i)
forall a b. (a -> b) -> a -> b
$ (Exp i -> Exp i -> Exp i)
-> Acc (Segments i)
-> (Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Segments i)
-> Acc (Segments i)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Acc (Segments i)
zeros (\Exp (DIM0 :. Int)
ix -> Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp i -> Exp (DIM0 :. Int)
forall i.
(Integral i, FromIntegral i Int) =>
Exp i -> Exp (DIM0 :. Int)
index1' (Acc (Scalar i) -> Exp i
forall e. Elt e => Acc (Scalar e) -> Exp e
the Acc (Scalar i)
len Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
- Exp i
1 Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
- Acc (Segments i)
offset Acc (Segments i) -> Exp (DIM0 :. Int) -> Exp i
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp (DIM0 :. Int)
ix))) Acc (Segments i)
ones
  where
    T2 Acc (Segments i)
offset Acc (Scalar i)
len = (Exp i -> Exp i -> Exp i)
-> Exp i -> Acc (Segments i) -> Acc (Segments i, Scalar i)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanr' Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
(+) Exp i
0 Acc (Segments i)
seg
    zeros :: Acc (Segments i)
zeros         = Exp (DIM0 :. Int) -> Exp i -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp i -> Exp (DIM0 :. Int)
forall i.
(Integral i, FromIntegral i Int) =>
Exp i -> Exp (DIM0 :. Int)
index1' (Exp i -> Exp (DIM0 :. Int)) -> Exp i -> Exp (DIM0 :. Int)
forall a b. (a -> b) -> a -> b
$ Acc (Scalar i) -> Exp i
forall e. Elt e => Acc (Scalar e) -> Exp e
the Acc (Scalar i)
len Exp i -> Exp i -> Exp i
forall a. Num a => a -> a -> a
+ Exp i
1) Exp i
0
    ones :: Acc (Segments i)
ones          = Exp (DIM0 :. Int) -> Exp i -> Acc (Segments i)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
index1  (Exp Int -> Exp (DIM0 :. Int)) -> Exp Int -> Exp (DIM0 :. Int)
forall a b. (a -> b) -> a -> b
$ Acc (Segments i) -> Exp Int
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
size Acc (Segments i)
offset) Exp i
1

-- | Construct a segmented version of a function from a non-segmented
-- version. The segmented apply operates on a head-flag value tuple, and
-- follows the procedure of Sengupta et. al.
--
segmentedL
    :: (Elt e, Num i, Bits i)
    => (Exp e -> Exp e -> Exp e)
    -> (Exp (i, e) -> Exp (i, e) -> Exp (i, e))
segmentedL :: (Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
segmentedL Exp e -> Exp e -> Exp e
f (T2 Exp i
aF Exp e
aV) (T2 Exp i
bF Exp e
bV) =
  Exp i -> Exp e -> Exp (i, e)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 (Exp i
aF Exp i -> Exp i -> Exp i
forall a. Bits a => Exp a -> Exp a -> Exp a
.|. Exp i
bF)
     (Exp i
bF Exp i -> Exp i -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
/= Exp i
0 Exp Bool -> (Exp e, Exp e) -> Exp e
forall t. Elt t => Exp Bool -> (Exp t, Exp t) -> Exp t
? (Exp e
bV, Exp e -> Exp e -> Exp e
f Exp e
aV Exp e
bV))

segmentedR
    :: (Elt e, Num i, Bits i)
    => (Exp e -> Exp e -> Exp e)
    -> (Exp (i, e) -> Exp (i, e) -> Exp (i, e))
segmentedR :: (Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
segmentedR Exp e -> Exp e -> Exp e
f Exp (i, e)
y Exp (i, e)
x = (Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
forall e i.
(Elt e, Num i, Bits i) =>
(Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
segmentedL ((Exp e -> Exp e -> Exp e) -> Exp e -> Exp e -> Exp e
forall a b c. (a -> b -> c) -> b -> a -> c
flip Exp e -> Exp e -> Exp e
f) Exp (i, e)
x Exp (i, e)
y

-- | Index construction and destruction generalised to integral types.
--
-- We generalise the segment descriptor to integral types because some
-- architectures, such as GPUs, have poor performance for 64-bit types. So,
-- there is a tension between performance and requiring 64-bit indices for some
-- applications, and we would not like to restrict ourselves to either one.
--
-- As we don't yet support non-Int dimensions in shapes, we will need to convert
-- back to concrete Int. However, don't put these generalised forms into the
-- base library, because it results in too many ambiguity errors.
--
index1' ::  (Integral i, FromIntegral i Int) => Exp i -> Exp DIM1
index1' :: Exp i -> Exp (DIM0 :. Int)
index1' Exp i
i = (DIM0 :. Exp Int) -> Exp (Plain (DIM0 :. Exp Int))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (DIM0
Z DIM0 -> Exp Int -> DIM0 :. Exp Int
forall tail head. tail -> head -> tail :. head
:. Exp i -> Exp Int
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral Exp i
i)


-- Reshaping of arrays
-- -------------------

-- | Flatten the given array of arbitrary dimension into a one-dimensional
-- vector. As with 'reshape', this operation performs no work.
--
flatten :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Acc (Vector e)
flatten :: Acc (Array sh e) -> Acc (Vector e)
flatten Acc (Array sh e)
a
  | Just sh :~: (DIM0 :. Int)
Refl <- (Shape sh, Shape (DIM0 :. Int)) => Maybe (sh :~: (DIM0 :. Int))
forall s t. (Shape s, Shape t) => Maybe (s :~: t)
matchShapeType @sh @DIM1
  = Acc (Array sh e)
Acc (Vector e)
a
flatten Acc (Array sh e)
a
  = Exp (DIM0 :. Int) -> Acc (Array sh e) -> Acc (Vector e)
forall sh sh' e.
(Shape sh, Shape sh', Elt e) =>
Exp sh -> Acc (Array sh' e) -> Acc (Array sh e)
reshape (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Array sh e) -> Exp Int
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
size Acc (Array sh e)
a)) Acc (Array sh e)
a


-- Enumeration and filling
-- -----------------------

-- | Create an array where all elements are the same value.
--
-- >>> run $ fill (constant (Z:.10)) 0 :: Vector Float
-- Vector (Z :. 10) [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]
--
fill :: (Shape sh, Elt e) => Exp sh -> Exp e -> Acc (Array sh e)
fill :: Exp sh -> Exp e -> Acc (Array sh e)
fill Exp sh
sh Exp e
c = Exp sh -> (Exp sh -> Exp e) -> Acc (Array sh e)
forall sh a.
(Shape sh, Elt a) =>
Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a)
generate Exp sh
sh (Exp e -> Exp sh -> Exp e
forall a b. a -> b -> a
const Exp e
c)

-- | Create an array of the given shape containing the values @x@, @x+1@, etc.
-- (in row-major order).
--
-- >>> run $ enumFromN (constant (Z:.5:.10)) 0 :: Matrix Int
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
enumFromN
    :: (Shape sh, Num e, FromIntegral Int e)
    => Exp sh
    -> Exp e
    -> Acc (Array sh e)
enumFromN :: Exp sh -> Exp e -> Acc (Array sh e)
enumFromN Exp sh
sh Exp e
x = Exp sh -> Exp e -> Exp e -> Acc (Array sh e)
forall sh e.
(Shape sh, Num e, FromIntegral Int e) =>
Exp sh -> Exp e -> Exp e -> Acc (Array sh e)
enumFromStepN Exp sh
sh Exp e
x Exp e
1

-- | Create an array of the given shape containing the values @x@, @x+y@,
-- @x+y+y@ etc. (in row-major order).
--
-- >>> run $ enumFromStepN (constant (Z:.5:.10)) 0 0.5 :: Matrix Float
-- Matrix (Z :. 5 :. 10)
--   [  0.0,  0.5,  1.0,  1.5,  2.0,  2.5,  3.0,  3.5,  4.0,  4.5,
--      5.0,  5.5,  6.0,  6.5,  7.0,  7.5,  8.0,  8.5,  9.0,  9.5,
--     10.0, 10.5, 11.0, 11.5, 12.0, 12.5, 13.0, 13.5, 14.0, 14.5,
--     15.0, 15.5, 16.0, 16.5, 17.0, 17.5, 18.0, 18.5, 19.0, 19.5,
--     20.0, 20.5, 21.0, 21.5, 22.0, 22.5, 23.0, 23.5, 24.0, 24.5]
--
enumFromStepN
    :: (Shape sh, Num e, FromIntegral Int e)
    => Exp sh
    -> Exp e              -- ^ x: start
    -> Exp e              -- ^ y: step
    -> Acc (Array sh e)
enumFromStepN :: Exp sh -> Exp e -> Exp e -> Acc (Array sh e)
enumFromStepN Exp sh
sh Exp e
x Exp e
y
  = Exp sh -> Acc (Array (DIM0 :. Int) e) -> Acc (Array sh e)
forall sh sh' e.
(Shape sh, Shape sh', Elt e) =>
Exp sh -> Acc (Array sh' e) -> Acc (Array sh e)
reshape Exp sh
sh
  (Acc (Array (DIM0 :. Int) e) -> Acc (Array sh e))
-> Acc (Array (DIM0 :. Int) e) -> Acc (Array sh e)
forall a b. (a -> b) -> a -> b
$ Exp (DIM0 :. Int)
-> (Exp (DIM0 :. Int) -> Exp e) -> Acc (Array (DIM0 :. Int) e)
forall sh a.
(Shape sh, Elt a) =>
Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a)
generate (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Exp sh -> Exp Int
forall sh. Shape sh => Exp sh -> Exp Int
shapeSize Exp sh
sh))
             (\Exp (DIM0 :. Int)
ix -> (Exp Int -> Exp e
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
fromIntegral (Exp (DIM0 :. Int) -> Exp Int
forall i. Elt i => Exp (DIM0 :. i) -> Exp i
unindex1 Exp (DIM0 :. Int)
ix :: Exp Int) Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
* Exp e
y) Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
+ Exp e
x)

-- Concatenation
-- -------------

-- | Concatenate innermost component of two arrays. The extent of the lower
--   dimensional component is the intersection of the two arrays.
--
-- >>> let m1 = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> m1
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> let m2 = fromList (Z:.10:.3) [0..] :: Matrix Int
-- >>> m2
-- Matrix (Z :. 10 :. 3)
--   [  0,  1,  2,
--      3,  4,  5,
--      6,  7,  8,
--      9, 10, 11,
--     12, 13, 14,
--     15, 16, 17,
--     18, 19, 20,
--     21, 22, 23,
--     24, 25, 26,
--     27, 28, 29]
--
-- >>> run $ use m1 ++ use m2
-- Matrix (Z :. 5 :. 13)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  0,  1,  2,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,  3,  4,  5,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,  6,  7,  8,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,  9, 10, 11,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 12, 13, 14]
--
infixr 5 ++
(++) :: (Shape sh, Elt e)
     => Acc (Array (sh :. Int) e)
     -> Acc (Array (sh :. Int) e)
     -> Acc (Array (sh :. Int) e)
++ :: Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
(++) = Lens' (Exp (sh :. Int)) (Exp Int)
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int)
-> Acc (Array sh e) -> Acc (Array sh e) -> Acc (Array sh e)
concatOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp (sh :. Int)) (Exp Int)
_1


-- | Generalised version of '(++)' where the argument 'Lens'' specifies which
-- dimension to concatenate along.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- >>> let m1 = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> m1
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> let m2 = fromList (Z:.10:.5) [0..] :: Matrix Int
-- >>> m2
-- Matrix (Z :. 10 :. 5)
--   [  0,  1,  2,  3,  4,
--      5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14,
--     15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24,
--     25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34,
--     35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44,
--     45, 46, 47, 48, 49]
--
-- >>> run $ concatOn _1 (use m1) (use m2)
-- Matrix (Z :. 5 :. 15)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  0,  1,  2,  3,  4,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,  5,  6,  7,  8,  9,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 10, 11, 12, 13, 14,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 15, 16, 17, 18, 19,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 20, 21, 22, 23, 24]
--
-- >>> run $ concatOn _2 (use m1) (use m2)
-- Matrix (Z :. 15 :. 5)
--   [  0,  1,  2,  3,  4,
--     10, 11, 12, 13, 14,
--     20, 21, 22, 23, 24,
--     30, 31, 32, 33, 34,
--     40, 41, 42, 43, 44,
--      0,  1,  2,  3,  4,
--      5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14,
--     15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24,
--     25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34,
--     35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44,
--     45, 46, 47, 48, 49]
--
concatOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Acc (Array sh e)
    -> Acc (Array sh e)
    -> Acc (Array sh e)
concatOn :: Lens' (Exp sh) (Exp Int)
-> Acc (Array sh e) -> Acc (Array sh e) -> Acc (Array sh e)
concatOn Lens' (Exp sh) (Exp Int)
dim Acc (Array sh e)
xs Acc (Array sh e)
ys =
  let
      shx :: Exp sh
shx   = Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
xs
      shy :: Exp sh
shy   = Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
ys
      m :: Exp Int
m     = Exp sh
shx Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim
      n :: Exp Int
n     = Exp sh
shy Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim
      shx' :: Exp sh
shx'  = Exp sh
shx Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Exp Int
mExp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+Exp Int
n
      shy' :: Exp sh
shy'  = Exp sh
shy Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Exp Int
mExp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+Exp Int
n
  in
  Exp sh -> (Exp sh -> Exp e) -> Acc (Array sh e)
forall sh a.
(Shape sh, Elt a) =>
Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a)
generate (Exp sh
shx' Exp sh -> Exp sh -> Exp sh
forall sh. Shape sh => Exp sh -> Exp sh -> Exp sh
`intersect` Exp sh
shy')
           (\Exp sh
ix -> let i :: Exp Int
i = Exp sh
ix Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim in
                   if  Exp Int
i Exp Int -> Exp Int -> Exp Bool
forall a. Ord a => Exp a -> Exp a -> Exp Bool
< Exp Int
m then Acc (Array sh e)
xs Acc (Array sh e) -> Exp sh -> Exp e
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp sh
ix
                             else Acc (Array sh e)
ys Acc (Array sh e) -> Exp sh -> Exp e
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! (Exp sh
ix Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall a s t. Num a => ASetter s t a a -> a -> s -> t
-~ Exp Int
m))

-- TLM: If we have something like (concat . split) then the source array will
--      have two use sites, but is actually safe (and better) to inline.


-- Filtering
-- ---------

-- | Drop elements that do not satisfy the predicate. Returns the elements which
-- pass the predicate, together with a segment descriptor indicating how many
-- elements along each outer dimension were valid.
--
-- >>> let vec = fromList (Z :. 10) [1..10] :: Vector Int
-- >>> vec
-- Vector (Z :. 10) [1,2,3,4,5,6,7,8,9,10]
--
-- >>> run $ filter even (use vec)
-- (Vector (Z :. 5) [2,4,6,8,10],Scalar Z [5])
--
-- >>> let mat = fromList (Z :. 4 :. 10) [1,2,3,4,5,6,7,8,9,10,1,1,1,1,1,2,2,2,2,2,2,4,6,8,10,12,14,16,18,20,1,3,5,7,9,11,13,15,17,19] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 4 :. 10)
--   [ 1, 2, 3, 4,  5,  6,  7,  8,  9, 10,
--     1, 1, 1, 1,  1,  2,  2,  2,  2,  2,
--     2, 4, 6, 8, 10, 12, 14, 16, 18, 20,
--     1, 3, 5, 7,  9, 11, 13, 15, 17, 19]
--
-- >>> run $ filter odd (use mat)
-- (Vector (Z :. 20) [1,3,5,7,9,1,1,1,1,1,1,3,5,7,9,11,13,15,17,19],Vector (Z :. 4) [5,5,0,10])
--
filter :: (Shape sh, Elt e)
       => (Exp e -> Exp Bool)
       -> Acc (Array (sh:.Int) e)
       -> Acc (Vector e, Array sh Int)
filter :: (Exp e -> Exp Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Vector e, Array sh Int)
filter Exp e -> Exp Bool
p Acc (Array (sh :. Int) e)
arr = Acc (Array (sh :. Int) Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Vector e, Array sh Int)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array (sh :. Int) Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Vector e, Array sh Int)
compact ((Exp e -> Exp Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) Bool)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp e -> Exp Bool
p Acc (Array (sh :. Int) e)
arr) Acc (Array (sh :. Int) e)
arr
{-# NOINLINE filter #-}
{-# RULES
  "ACC filter/filter" forall f g arr.
    filter f (afst (filter g arr)) = filter (\x -> g x && f x) arr
 #-}


-- | As 'filter', but with separate arrays for the data elements and the
-- flags indicating which elements of that array should be kept.
--
compact :: forall sh e. (Shape sh, Elt e)
        => Acc (Array (sh:.Int) Bool)
        -> Acc (Array (sh:.Int) e)
        -> Acc (Vector e, Array sh Int)
compact :: Acc (Array (sh :. Int) Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Vector e, Array sh Int)
compact Acc (Array (sh :. Int) Bool)
keep Acc (Array (sh :. Int) e)
arr
  -- Optimise 1-dimensional arrays, where we can avoid additional computations
  -- for the offset indices.
  | Just sh :~: DIM0
Refl <- (Shape sh, Shape DIM0) => Maybe (sh :~: DIM0)
forall s t. (Shape s, Shape t) => Maybe (s :~: t)
matchShapeType @sh @Z
  = let
        T2 Acc (Array (sh :. Int) Int)
target Acc (Array sh Int)
len   = (Exp Int -> Exp Int -> Exp Int)
-> Exp Int
-> Acc (Array (sh :. Int) Int)
-> Acc (Array (sh :. Int) Int, Array sh Int)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanl' Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
(+) Exp Int
0 ((Exp Bool -> Exp Int)
-> Acc (Array (sh :. Int) Bool) -> Acc (Array (sh :. Int) Int)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp Bool -> Exp Int
boolToInt Acc (Array (sh :. Int) Bool)
keep)
        prj :: Exp (sh :. Int) -> Exp (Maybe (DIM0 :. Int))
prj Exp (sh :. Int)
ix          = if Acc (Array (sh :. Int) Bool)
keepAcc (Array (sh :. Int) Bool) -> Exp (sh :. Int) -> Exp Bool
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
!Exp (sh :. Int)
ix
                             then Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Array (sh :. Int) Int)
targetAcc (Array (sh :. Int) Int) -> Exp (sh :. Int) -> Exp Int
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
!Exp (sh :. Int)
ix))
                             else Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp (Maybe a)
Nothing_
        dummy :: Acc (Vector e)
dummy           = Exp (DIM0 :. Int) -> Exp e -> Acc (Vector e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Scalar Int) -> Exp Int
forall e. Elt e => Acc (Scalar e) -> Exp e
the Acc (Array sh Int)
Acc (Scalar Int)
len)) Exp e
forall e. Elt e => Exp e
undef
        result :: Acc (Vector e)
result          = (Exp e -> Exp e -> Exp e)
-> Acc (Vector e)
-> (Exp (sh :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Array (sh :. Int) e)
-> Acc (Vector e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp e -> Exp e -> Exp e
forall a b. a -> b -> a
const Acc (Vector e)
dummy Exp (sh :. Int) -> Exp (Maybe (DIM0 :. Int))
prj Acc (Array (sh :. Int) e)
arr
    in
    if Acc (Array (sh :. Int) e) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Array (sh :. Int) e)
arr
      then Acc (Vector e) -> Acc (Scalar Int) -> Acc (Vector e, Scalar Int)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Vector e)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e)
emptyArray (Exp DIM0 -> Exp Int -> Acc (Scalar Int)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill Exp DIM0
Z_ Exp Int
0)
      else Acc (Vector e)
-> Acc (Array sh Int) -> Acc (Vector e, Array sh Int)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Vector e)
result Acc (Array sh Int)
len

compact Acc (Array (sh :. Int) Bool)
keep Acc (Array (sh :. Int) e)
arr
  = let
        sz :: Exp sh
sz              = Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail (Acc (Array (sh :. Int) e) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) e)
arr)
        T2 Acc (Array (sh :. Int) Int)
target Acc (Array sh Int)
len   = (Exp Int -> Exp Int -> Exp Int)
-> Exp Int
-> Acc (Array (sh :. Int) Int)
-> Acc (Array (sh :. Int) Int, Array sh Int)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanl' Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
(+) Exp Int
0 ((Exp Bool -> Exp Int)
-> Acc (Array (sh :. Int) Bool) -> Acc (Array (sh :. Int) Int)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp Bool -> Exp Int
boolToInt Acc (Array (sh :. Int) Bool)
keep)
        T2 Acc (Array (DIM0 :. Int) Int)
offset Acc (Scalar Int)
valid = (Exp Int -> Exp Int -> Exp Int)
-> Exp Int
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) Int, Scalar Int)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanl' Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
(+) Exp Int
0 (Acc (Array sh Int) -> Acc (Array (DIM0 :. Int) Int)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Acc (Vector e)
flatten Acc (Array sh Int)
len)
        prj :: Exp (sh :. Int) -> Exp (Maybe (DIM0 :. Int))
prj Exp (sh :. Int)
ix          = if Acc (Array (sh :. Int) Bool)
keepAcc (Array (sh :. Int) Bool) -> Exp (sh :. Int) -> Exp Bool
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
!Exp (sh :. Int)
ix
                            then Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Array (DIM0 :. Int) Int)
offset Acc (Array (DIM0 :. Int) Int) -> Exp Int -> Exp Int
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp Int -> Exp e
!! (Exp sh -> Exp sh -> Exp Int
forall sh. Shape sh => Exp sh -> Exp sh -> Exp Int
toIndex Exp sh
sz (Exp (sh :. Int) -> Exp sh
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp sh
indexTail Exp (sh :. Int)
ix)) Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Acc (Array (sh :. Int) Int)
targetAcc (Array (sh :. Int) Int) -> Exp (sh :. Int) -> Exp Int
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
!Exp (sh :. Int)
ix))
                            else Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp (Maybe a)
Nothing_
        dummy :: Acc (Vector e)
dummy           = Exp (DIM0 :. Int) -> Exp e -> Acc (Vector e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Scalar Int) -> Exp Int
forall e. Elt e => Acc (Scalar e) -> Exp e
the Acc (Scalar Int)
valid)) Exp e
forall e. Elt e => Exp e
undef
        result :: Acc (Vector e)
result          = (Exp e -> Exp e -> Exp e)
-> Acc (Vector e)
-> (Exp (sh :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Array (sh :. Int) e)
-> Acc (Vector e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp e -> Exp e -> Exp e
forall a b. a -> b -> a
const Acc (Vector e)
dummy Exp (sh :. Int) -> Exp (Maybe (DIM0 :. Int))
prj Acc (Array (sh :. Int) e)
arr
    in
    if Acc (Array (sh :. Int) e) -> Exp Bool
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null Acc (Array (sh :. Int) e)
arr
      then Acc (Vector e)
-> Acc (Array sh Int) -> Acc (Vector e, Array sh Int)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Vector e)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e)
emptyArray (Exp sh -> Exp Int -> Acc (Array sh Int)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill Exp sh
sz Exp Int
0)
      else Acc (Vector e)
-> Acc (Array sh Int) -> Acc (Vector e, Array sh Int)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Acc (Vector e)
result Acc (Array sh Int)
len


-- Gather operations
-- -----------------

-- | Gather elements from a source array by reading values at the given indices.
--
-- >>> let input = fromList (Z:.9) [1,9,6,4,4,2,0,1,2] :: Vector Int
-- >>> let from  = fromList (Z:.6) [1,3,7,2,5,3] :: Vector Int
-- >>> run $ gather (use from) (use input)
-- Vector (Z :. 6) [9,4,1,6,2,4]
--
gather
    :: (Shape sh, Elt e)
    => Acc (Array sh Int)         -- ^ index of source at each index to gather
    -> Acc (Vector e)             -- ^ source values
    -> Acc (Array sh e)
gather :: Acc (Array sh Int) -> Acc (Vector e) -> Acc (Array sh e)
gather Acc (Array sh Int)
indices Acc (Vector e)
input = (Exp Int -> Exp e) -> Acc (Array sh Int) -> Acc (Array sh e)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Acc (Vector e)
input Acc (Vector e) -> Exp Int -> Exp e
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp Int -> Exp e
!!) Acc (Array sh Int)
indices
  -- TLM NOTES:
  --  * (!!) has potential for later optimisation
  --  * We needn't fix the source array to Vector, but this matches the
  --    intuition that 'Int' ~ 'DIM1'.


-- | Conditionally copy elements from source array to destination array
-- according to an index mapping.
--
-- In addition, the 'mask' vector and associated predication function specifies
-- whether the element is copied or a default value is used instead.
--
-- >>> let defaults = fromList (Z :. 6) [6,6,6,6,6,6] :: Vector Float
-- >>> let from     = fromList (Z :. 6) [1,3,7,2,5,3] :: Vector Int
-- >>> let mask     = fromList (Z :. 6) [3,4,9,2,7,5] :: Vector Int
-- >>> let input    = fromList (Z :. 9) [1,9,6,4,4,2,0,1,2] :: Vector Float
-- >>> run $ gatherIf (use from) (use mask) (> 4) (use defaults) (use input)
-- Vector (Z :. 6) [6.0,6.0,1.0,6.0,2.0,4.0]
--
gatherIf
    :: (Elt a, Elt b)
    => Acc (Vector Int)           -- ^ source indices to gather from
    -> Acc (Vector a)             -- ^ mask vector
    -> (Exp a -> Exp Bool)        -- ^ predicate function
    -> Acc (Vector b)             -- ^ default values
    -> Acc (Vector b)             -- ^ source values
    -> Acc (Vector b)
gatherIf :: Acc (Array (DIM0 :. Int) Int)
-> Acc (Vector a)
-> (Exp a -> Exp Bool)
-> Acc (Vector b)
-> Acc (Vector b)
-> Acc (Vector b)
gatherIf Acc (Array (DIM0 :. Int) Int)
from Acc (Vector a)
maskV Exp a -> Exp Bool
pred Acc (Vector b)
defaults Acc (Vector b)
input = (Exp Bool -> Exp (b, b) -> Exp b)
-> Acc (Array (DIM0 :. Int) Bool)
-> Acc (Array (DIM0 :. Int) (b, b))
-> Acc (Vector b)
forall sh a b c.
(Shape sh, Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
zipWith Exp Bool -> Exp (b, b) -> Exp b
forall t (c :: * -> *).
(Elt t, Unlift c (Exp t, Exp t)) =>
Exp Bool -> c (t, t) -> Exp t
zf Acc (Array (DIM0 :. Int) Bool)
pf Acc (Array (DIM0 :. Int) (b, b))
gatheredV
  where
    zf :: Exp Bool -> c (t, t) -> Exp t
zf Exp Bool
p c (t, t)
g      = Exp Bool
p Exp Bool -> (Exp t, Exp t) -> Exp t
forall t. Elt t => Exp Bool -> (Exp t, Exp t) -> Exp t
? (c (Plain (Exp t, Exp t)) -> (Exp t, Exp t)
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift c (t, t)
c (Plain (Exp t, Exp t))
g)
    gatheredV :: Acc (Array (DIM0 :. Int) (b, b))
gatheredV   = Acc (Vector b)
-> Acc (Vector b) -> Acc (Array (DIM0 :. Int) (b, b))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip (Acc (Array (DIM0 :. Int) Int) -> Acc (Vector b) -> Acc (Vector b)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh Int) -> Acc (Vector e) -> Acc (Array sh e)
gather Acc (Array (DIM0 :. Int) Int)
from Acc (Vector b)
input) Acc (Vector b)
defaults
    pf :: Acc (Array (DIM0 :. Int) Bool)
pf          = (Exp a -> Exp Bool)
-> Acc (Vector a) -> Acc (Array (DIM0 :. Int) Bool)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp a -> Exp Bool
pred Acc (Vector a)
maskV


-- Scatter operations
-- ------------------

-- | Overwrite elements of the destination by scattering the values of the
-- source array according to the given index mapping.
--
-- Note that if the destination index appears more than once in the mapping the
-- result is undefined.
--
-- >>> let to    = fromList (Z :. 6) [1,3,7,2,5,8] :: Vector Int
-- >>> let input = fromList (Z :. 7) [1,9,6,4,4,2,5] :: Vector Int
-- >>> run $ scatter (use to) (fill (constant (Z:.10)) 0) (use input)
-- Vector (Z :. 10) [0,1,4,9,0,4,0,6,2,0]
--
scatter
    :: Elt e
    => Acc (Vector Int)           -- ^ destination indices to scatter into
    -> Acc (Vector e)             -- ^ default values
    -> Acc (Vector e)             -- ^ source values
    -> Acc (Vector e)
scatter :: Acc (Array (DIM0 :. Int) Int)
-> Acc (Vector e) -> Acc (Vector e) -> Acc (Vector e)
scatter Acc (Array (DIM0 :. Int) Int)
to Acc (Vector e)
defaults Acc (Vector e)
input = (Exp e -> Exp e -> Exp e)
-> Acc (Vector e)
-> (Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Vector e)
-> Acc (Vector e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp e -> Exp e -> Exp e
forall a b. a -> b -> a
const Acc (Vector e)
defaults Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
pf Acc (Vector e)
input'
  where
    pf :: Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
pf Exp (DIM0 :. Int)
ix   = Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Array (DIM0 :. Int) Int)
to Acc (Array (DIM0 :. Int) Int) -> Exp (DIM0 :. Int) -> Exp Int
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp (DIM0 :. Int)
ix))
    input' :: Acc (Vector e)
input'  = Exp (DIM0 :. Int)
-> (Exp (DIM0 :. Int) -> Exp (DIM0 :. Int))
-> Acc (Vector e)
-> Acc (Vector e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute (Acc (Array (DIM0 :. Int) Int) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (DIM0 :. Int) Int)
to Exp (DIM0 :. Int) -> Exp (DIM0 :. Int) -> Exp (DIM0 :. Int)
forall sh. Shape sh => Exp sh -> Exp sh -> Exp sh
`intersect` Acc (Vector e) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Vector e)
input) Exp (DIM0 :. Int) -> Exp (DIM0 :. Int)
forall a. a -> a
id Acc (Vector e)
input


-- | Conditionally overwrite elements of the destination by scattering values of
-- the source array according to a given index mapping, whenever the masking
-- function resolves to 'True'.
--
-- Note that if the destination index appears more than once in the mapping the
-- result is undefined.
--
-- >>> let to    = fromList (Z :. 6) [1,3,7,2,5,8] :: Vector Int
-- >>> let mask  = fromList (Z :. 6) [3,4,9,2,7,5] :: Vector Int
-- >>> let input = fromList (Z :. 7) [1,9,6,4,4,2,5] :: Vector Int
-- >>> run $ scatterIf (use to) (use mask) (> 4) (fill (constant (Z:.10)) 0) (use input)
-- Vector (Z :. 10) [0,0,0,0,0,4,0,6,2,0]
--
scatterIf
    :: (Elt a, Elt b)
    => Acc (Vector Int)           -- ^ destination indices to scatter into
    -> Acc (Vector a)             -- ^ mask vector
    -> (Exp a -> Exp Bool)        -- ^ predicate function
    -> Acc (Vector b)             -- ^ default values
    -> Acc (Vector b)             -- ^ source values
    -> Acc (Vector b)
scatterIf :: Acc (Array (DIM0 :. Int) Int)
-> Acc (Vector a)
-> (Exp a -> Exp Bool)
-> Acc (Vector b)
-> Acc (Vector b)
-> Acc (Vector b)
scatterIf Acc (Array (DIM0 :. Int) Int)
to Acc (Vector a)
maskV Exp a -> Exp Bool
pred Acc (Vector b)
defaults Acc (Vector b)
input = (Exp b -> Exp b -> Exp b)
-> Acc (Vector b)
-> (Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Vector b)
-> Acc (Vector b)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp b -> Exp b -> Exp b
forall a b. a -> b -> a
const Acc (Vector b)
defaults Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
pf Acc (Vector b)
input'
  where
    input' :: Acc (Vector b)
input'  = Exp (DIM0 :. Int)
-> (Exp (DIM0 :. Int) -> Exp (DIM0 :. Int))
-> Acc (Vector b)
-> Acc (Vector b)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute (Acc (Array (DIM0 :. Int) Int) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (DIM0 :. Int) Int)
to Exp (DIM0 :. Int) -> Exp (DIM0 :. Int) -> Exp (DIM0 :. Int)
forall sh. Shape sh => Exp sh -> Exp sh -> Exp sh
`intersect` Acc (Vector b) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Vector b)
input) Exp (DIM0 :. Int) -> Exp (DIM0 :. Int)
forall a. a -> a
id Acc (Vector b)
input
    pf :: Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
pf Exp (DIM0 :. Int)
ix   = if Exp a -> Exp Bool
pred (Acc (Vector a)
maskV Acc (Vector a) -> Exp (DIM0 :. Int) -> Exp a
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp (DIM0 :. Int)
ix)
                 then Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Array (DIM0 :. Int) Int)
to Acc (Array (DIM0 :. Int) Int) -> Exp (DIM0 :. Int) -> Exp Int
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp (DIM0 :. Int)
ix))
                 else Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp (Maybe a)
Nothing_


-- Permutations
-- ------------

-- | Reverse the elements of a vector.
--
reverse :: Elt e => Acc (Vector e) -> Acc (Vector e)
reverse :: Acc (Vector e) -> Acc (Vector e)
reverse = Lens' (Exp (DIM0 :. Int)) (Exp Int)
-> Acc (Vector e) -> Acc (Vector e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
reverseOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp (DIM0 :. Int)) (Exp Int)
_1

-- | Transpose the rows and columns of a matrix.
--
transpose :: Elt e => Acc (Array DIM2 e) -> Acc (Array DIM2 e)
transpose :: Acc (Array DIM2 e) -> Acc (Array DIM2 e)
transpose = Lens' (Exp DIM2) (Exp Int)
-> Lens' (Exp DIM2) (Exp Int)
-> Acc (Array DIM2 e)
-> Acc (Array DIM2 e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int)
-> Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
transposeOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp DIM2) (Exp Int)
_1 forall sh. Elt sh => Lens' (Exp ((sh :. Int) :. Int)) (Exp Int)
Lens' (Exp DIM2) (Exp Int)
_2


-- | Generalised version of 'reverse' where the argument 'Lens'' specifies which
-- dimension to reverse.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ reverseOn _1 (use mat)
-- Matrix (Z :. 5 :. 10)
--   [  9,  8,  7,  6,  5,  4,  3,  2,  1,  0,
--     19, 18, 17, 16, 15, 14, 13, 12, 11, 10,
--     29, 28, 27, 26, 25, 24, 23, 22, 21, 20,
--     39, 38, 37, 36, 35, 34, 33, 32, 31, 30,
--     49, 48, 47, 46, 45, 44, 43, 42, 41, 40]
--
-- >>> run $ reverseOn _2 (use mat)
-- Matrix (Z :. 5 :. 10)
--   [ 40, 41, 42, 43, 44, 45, 46, 47, 48, 49,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--      0,  1,  2,  3,  4,  5,  6,  7,  8,  9]
--
-- @since 1.2.0.0
--
reverseOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Acc (Array sh e)
    -> Acc (Array sh e)
reverseOn :: Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
reverseOn Lens' (Exp sh) (Exp Int)
dim Acc (Array sh e)
xs =
  let
      sh :: Exp sh
sh = Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
xs
      n :: Exp Int
n  = Exp sh
sh Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim
  in
  Exp sh
-> (Exp sh -> Exp sh) -> Acc (Array sh e) -> Acc (Array sh e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute Exp sh
sh (ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
-> (Exp Int -> Exp Int) -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim ((Exp Int -> Exp Int) -> Exp sh -> Exp sh)
-> (Exp Int -> Exp Int) -> Exp sh -> Exp sh
forall a b. (a -> b) -> a -> b
$ \Exp Int
i -> Exp Int
n Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
- Exp Int
i Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
- Exp Int
1) Acc (Array sh e)
xs

-- | Generalised version of 'transpose' where the argument 'Lens''s specify
-- which two dimensions to transpose.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ transposeOn _1 _2 (use mat)
-- Matrix (Z :. 10 :. 5)
--   [ 0, 10, 20, 30, 40,
--     1, 11, 21, 31, 41,
--     2, 12, 22, 32, 42,
--     3, 13, 23, 33, 43,
--     4, 14, 24, 34, 44,
--     5, 15, 25, 35, 45,
--     6, 16, 26, 36, 46,
--     7, 17, 27, 37, 47,
--     8, 18, 28, 38, 48,
--     9, 19, 29, 39, 49]
--
-- >>> let box = fromList (Z:.2:.3:.5) [0..] :: Array DIM3 Int
-- >>> box
-- Array (Z :. 2 :. 3 :. 5) [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29]
--
-- >>> run $ transposeOn _1 _2 (use box)
-- Array (Z :. 2 :. 5 :. 3) [0,5,10,1,6,11,2,7,12,3,8,13,4,9,14,15,20,25,16,21,26,17,22,27,18,23,28,19,24,29]
--
-- >>> run $ transposeOn _2 _3 (use box)
-- Array (Z :. 3 :. 2 :. 5) [0,1,2,3,4,15,16,17,18,19,5,6,7,8,9,20,21,22,23,24,10,11,12,13,14,25,26,27,28,29]
--
-- >>> run $ transposeOn _1 _3 (use box)
-- Array (Z :. 5 :. 3 :. 2) [0,15,5,20,10,25,1,16,6,21,11,26,2,17,7,22,12,27,3,18,8,23,13,28,4,19,9,24,14,29]
--
-- @since 1.2.0.0
--
transposeOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Lens' (Exp sh) (Exp Int)
    -> Acc (Array sh e)
    -> Acc (Array sh e)
transposeOn :: Lens' (Exp sh) (Exp Int)
-> Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
transposeOn Lens' (Exp sh) (Exp Int)
dim1 Lens' (Exp sh) (Exp Int)
dim2 Acc (Array sh e)
xs =
  let
      swap :: Exp sh -> Exp sh
swap Exp sh
ix = Exp sh
ix Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim2 ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Exp sh
ix Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim1
                   Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim1 ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Exp sh
ix Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim2
  in
  Exp sh
-> (Exp sh -> Exp sh) -> Acc (Array sh e) -> Acc (Array sh e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute (Exp sh -> Exp sh
swap (Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
xs)) Exp sh -> Exp sh
swap Acc (Array sh e)
xs


-- Extracting sub-vectors
-- ----------------------

-- | Yield the first @n@ elements in the innermost dimension of the array (plus
-- all lower dimensional elements).
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ take 5 (use mat)
-- Matrix (Z :. 5 :. 5)
--   [  0,  1,  2,  3,  4,
--     10, 11, 12, 13, 14,
--     20, 21, 22, 23, 24,
--     30, 31, 32, 33, 34,
--     40, 41, 42, 43, 44]
--
take :: (Shape sh, Elt e)
     => Exp Int
     -> Acc (Array (sh :. Int) e)
     -> Acc (Array (sh :. Int) e)
take :: Exp Int -> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
take = Lens' (Exp (sh :. Int)) (Exp Int)
-> Exp Int
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
takeOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp (sh :. Int)) (Exp Int)
_1


-- | Yield all but the first @n@ elements along the innermost dimension of the
-- array (plus all lower dimensional elements).
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ drop 7 (use mat)
-- Matrix (Z :. 5 :. 3)
--   [  7,  8,  9,
--     17, 18, 19,
--     27, 28, 29,
--     37, 38, 39,
--     47, 48, 49]
--
drop :: (Shape sh, Elt e)
     => Exp Int
     -> Acc (Array (sh :. Int) e)
     -> Acc (Array (sh :. Int) e)
drop :: Exp Int -> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
drop = Lens' (Exp (sh :. Int)) (Exp Int)
-> Exp Int
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
dropOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp (sh :. Int)) (Exp Int)
_1


-- | Yield all but the elements in the last index of the innermost dimension.
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ init (use mat)
-- Matrix (Z :. 5 :. 9)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,
--     10, 11, 12, 13, 14, 15, 16, 17, 18,
--     20, 21, 22, 23, 24, 25, 26, 27, 28,
--     30, 31, 32, 33, 34, 35, 36, 37, 38,
--     40, 41, 42, 43, 44, 45, 46, 47, 48]
--
init :: (Shape sh, Elt e)
     => Acc (Array (sh :. Int) e)
     -> Acc (Array (sh :. Int) e)
init :: Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
init = Lens' (Exp (sh :. Int)) (Exp Int)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
initOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp (sh :. Int)) (Exp Int)
_1


-- | Yield all but the first element along the innermost dimension of an array.
-- The innermost dimension must not be empty.
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ tail (use mat)
-- Matrix (Z :. 5 :. 9)
--   [  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     11, 12, 13, 14, 15, 16, 17, 18, 19,
--     21, 22, 23, 24, 25, 26, 27, 28, 29,
--     31, 32, 33, 34, 35, 36, 37, 38, 39,
--     41, 42, 43, 44, 45, 46, 47, 48, 49]
--
tail :: (Shape sh, Elt e)
     => Acc (Array (sh :. Int) e)
     -> Acc (Array (sh :. Int) e)
tail :: Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
tail = Lens' (Exp (sh :. Int)) (Exp Int)
-> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
tailOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp (sh :. Int)) (Exp Int)
_1


-- | Yield a slit (slice) of the innermost indices of an array. Denotationally,
-- we have:
--
-- > slit i n = take n . drop i
--
slit :: (Shape sh, Elt e)
     => Exp Int                     -- ^ starting index
     -> Exp Int                     -- ^ length
     -> Acc (Array (sh :. Int) e)
     -> Acc (Array (sh :. Int) e)
slit :: Exp Int
-> Exp Int
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
slit = Lens' (Exp (sh :. Int)) (Exp Int)
-> Exp Int
-> Exp Int
-> Acc (Array (sh :. Int) e)
-> Acc (Array (sh :. Int) e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int)
-> Exp Int -> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
slitOn forall sh. Elt sh => Lens' (Exp (sh :. Int)) (Exp Int)
Lens' (Exp (sh :. Int)) (Exp Int)
_1


-- | Generalised version of 'init' where the argument 'Lens'' specifies which
-- dimension to operate over.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- @since 1.2.0.0
--
initOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Acc (Array sh e)
    -> Acc (Array sh e)
initOn :: Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
initOn Lens' (Exp sh) (Exp Int)
dim Acc (Array sh e)
xs =
  let
      sh :: Exp sh
sh  = Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
xs
      sh' :: Exp sh
sh' = ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
-> (Exp Int -> Exp Int) -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim (\Exp Int
i -> Exp Int
0 Exp Int -> Exp Int -> Exp Int
forall a. Ord a => Exp a -> Exp a -> Exp a
`max` (Exp Int
iExp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
-Exp Int
1)) Exp sh
sh
  in
  Exp sh
-> (Exp sh -> Exp sh) -> Acc (Array sh e) -> Acc (Array sh e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute Exp sh
sh' Exp sh -> Exp sh
forall a. a -> a
id Acc (Array sh e)
xs


-- | Generalised version of 'tail' where the argument 'Lens'' specifies which
-- dimension to operate over.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- @since 1.2.0.0
--
tailOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Acc (Array sh e)
    -> Acc (Array sh e)
tailOn :: Lens' (Exp sh) (Exp Int) -> Acc (Array sh e) -> Acc (Array sh e)
tailOn Lens' (Exp sh) (Exp Int)
dim Acc (Array sh e)
xs =
  let
      sh :: Exp sh
sh  = Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
xs
      sh' :: Exp sh
sh' = ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
-> (Exp Int -> Exp Int) -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim (\Exp Int
i -> Exp Int
0 Exp Int -> Exp Int -> Exp Int
forall a. Ord a => Exp a -> Exp a -> Exp a
`max` (Exp Int
iExp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
-Exp Int
1)) Exp sh
sh
  in
  Exp sh
-> (Exp sh -> Exp sh) -> Acc (Array sh e) -> Acc (Array sh e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute Exp sh
sh' (Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim ASetter (Exp sh) (Exp sh) (Exp Int) (Exp Int)
-> Exp Int -> Exp sh -> Exp sh
forall a s t. Num a => ASetter s t a a -> a -> s -> t
+~ Exp Int
1) Acc (Array sh e)
xs


-- | Generalised version of 'take' where the argument 'Lens'' specifies which
-- dimension to operate over.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- @since 1.2.0.0
--
takeOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Exp Int
    -> Acc (Array sh e)
    -> Acc (Array sh e)
takeOn :: Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
takeOn Lens' (Exp sh) (Exp Int)
dim Exp Int
n Acc (Array sh e)
xs =
  let
      sh :: Exp sh
sh = Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
xs
      m :: Exp Int
m  = Exp sh
sh Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim
  in
  Exp sh
-> (Exp sh -> Exp sh) -> Acc (Array sh e) -> Acc (Array sh e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute (Exp sh
sh Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Exp Int -> Exp Int -> Exp Int
forall a. Ord a => Exp a -> Exp a -> Exp a
min Exp Int
m Exp Int
n) Exp sh -> Exp sh
forall a. a -> a
id Acc (Array sh e)
xs


-- | Generalised version of 'drop' where the argument 'Lens'' specifies which
-- dimension to operate over.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- @since 1.2.0.0
--
dropOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Exp Int
    -> Acc (Array sh e)
    -> Acc (Array sh e)
dropOn :: Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
dropOn Lens' (Exp sh) (Exp Int)
dim Exp Int
n Acc (Array sh e)
xs =
  let
      sh :: Exp sh
sh = Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array sh e)
xs
      m :: Exp Int
m  = Exp sh
sh Exp sh -> Getting (Exp Int) (Exp sh) (Exp Int) -> Exp Int
forall s a. s -> Getting a s a -> a
^. Getting (Exp Int) (Exp sh) (Exp Int)
Lens' (Exp sh) (Exp Int)
dim
  in
  Exp sh
-> (Exp sh -> Exp sh) -> Acc (Array sh e) -> Acc (Array sh e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute (Exp sh
sh Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Exp Int -> Exp Int -> Exp Int
forall a. Ord a => Exp a -> Exp a -> Exp a
max Exp Int
0 (Exp Int
mExp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
-Exp Int
n)) (Exp sh -> (Exp sh -> Exp sh) -> Exp sh
forall a b. a -> (a -> b) -> b
& (Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh)
Lens' (Exp sh) (Exp Int)
dim ((Exp Int -> Identity (Exp Int)) -> Exp sh -> Identity (Exp sh))
-> Exp Int -> Exp sh -> Exp sh
forall a s t. Num a => ASetter s t a a -> a -> s -> t
+~ Exp Int
n) Acc (Array sh e)
xs

-- Note: [embedding constants in take & drop]
--
-- Previously the 'take' and 'drop functions prevented the value of the
-- take/drop amount from being embedded directly in the generated code. This was
-- done by writing the value into a scalar array and reading that value out,
-- rather than using it directly. Although that is better from a code
-- cache/reuse standpoint, I've now removed this as it prevents the user from
-- specialising their code, and in a real program this extra indirection is
-- probably not necessary anyway.
--

-- | Generalised version of 'drop' where the argument 'Lens'' specifies which
-- dimension to operate over.
--
-- Appropriate lenses are available from <https://hackage.haskell.org/package/lens-accelerate lens-accelerate>.
--
-- @since 1.2.0.0
--
slitOn
    :: (Shape sh, Elt e)
    => Lens' (Exp sh) (Exp Int)
    -> Exp Int                      -- ^ starting index
    -> Exp Int                      -- ^ length
    -> Acc (Array sh e)
    -> Acc (Array sh e)
slitOn :: Lens' (Exp sh) (Exp Int)
-> Exp Int -> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
slitOn Lens' (Exp sh) (Exp Int)
dim Exp Int
i Exp Int
n = Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
takeOn Lens' (Exp sh) (Exp Int)
dim Exp Int
n (Acc (Array sh e) -> Acc (Array sh e))
-> (Acc (Array sh e) -> Acc (Array sh e))
-> Acc (Array sh e)
-> Acc (Array sh e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
forall sh e.
(Shape sh, Elt e) =>
Lens' (Exp sh) (Exp Int)
-> Exp Int -> Acc (Array sh e) -> Acc (Array sh e)
dropOn Lens' (Exp sh) (Exp Int)
dim Exp Int
i


-- Controlling execution
-- ---------------------

-- | Force an array expression to be evaluated, preventing it from fusing with
-- other operations. Forcing operations to be computed to memory, rather than
-- being fused into their consuming function, can sometimes improve performance.
-- For example, computing a matrix 'transpose' could provide better memory
-- locality for the subsequent operation. Preventing fusion to split large
-- operations into several simpler steps could also help by reducing register
-- pressure.
--
-- Preventing fusion also means that the individual operations are available to
-- be executed concurrently with other kernels. In particular, consider using
-- this if you have a series of operations that are compute bound rather than
-- memory bound.
--
-- Here is the synthetic example:
--
-- > loop :: Exp Int -> Exp Int
-- > loop ticks =
-- >   let clockRate = 900000   -- kHz
-- >   in  while (\i -> i < clockRate * ticks) (+1) 0
-- >
-- > test :: Acc (Vector Int)
-- > test =
-- >   zip3
-- >     (compute $ map loop (use $ fromList (Z:.1) [10]))
-- >     (compute $ map loop (use $ fromList (Z:.1) [10]))
-- >     (compute $ map loop (use $ fromList (Z:.1) [10]))
-- >
--
-- Without the use of 'compute', the operations are fused together and the three
-- long-running loops are executed sequentially in a single kernel. Instead, the
-- individual operations can now be executed concurrently, potentially reducing
-- overall runtime.
--
compute :: Arrays a => Acc a -> Acc a
compute :: Acc a -> Acc a
compute = Acc a -> Acc a
forall a. a -> a
id (Acc a -> Acc a) -> (Acc a -> Acc a) -> Acc a -> Acc a
forall a b c.
(Arrays a, Arrays b, Arrays c) =>
(Acc a -> Acc b) -> (Acc b -> Acc c) -> Acc a -> Acc c
>-> Acc a -> Acc a
forall a. a -> a
id


-- Flow control
-- ------------

-- | Infix version of 'acond'. If the predicate evaluates to 'True', the first
-- component of the tuple is returned, else the second.
--
-- Enabling the @RebindableSyntax@ extension will allow you to use the standard
-- if-then-else syntax instead.
--
infix 0 ?|
(?|) :: Arrays a => Exp Bool -> (Acc a, Acc a) -> Acc a
Exp Bool
c ?| :: Exp Bool -> (Acc a, Acc a) -> Acc a
?| (Acc a
t, Acc a
e) = Exp Bool -> Acc a -> Acc a -> Acc a
forall a. Arrays a => Exp Bool -> Acc a -> Acc a -> Acc a
acond Exp Bool
c Acc a
t Acc a
e

-- | An infix version of 'cond'. If the predicate evaluates to 'True', the first
-- component of the tuple is returned, else the second.
--
-- Enabling the @RebindableSyntax@ extension will allow you to use the standard
-- if-then-else syntax instead.
--
infix 0 ?
(?) :: Elt t => Exp Bool -> (Exp t, Exp t) -> Exp t
Exp Bool
c ? :: Exp Bool -> (Exp t, Exp t) -> Exp t
? (Exp t
t, Exp t
e) = Exp Bool -> Exp t -> Exp t -> Exp t
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond Exp Bool
c Exp t
t Exp t
e

-- | For use with @-XRebindableSyntax@, this class provides 'ifThenElse' lifted
-- to both scalar and array types.
--
class IfThenElse t where
  type EltT t a :: Constraint
  ifThenElse :: EltT t a => Exp Bool -> t a -> t a -> t a

instance IfThenElse Exp where
  type EltT Exp t = Elt t
  ifThenElse :: Exp Bool -> Exp a -> Exp a -> Exp a
ifThenElse = Exp Bool -> Exp a -> Exp a -> Exp a
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond

instance IfThenElse Acc where
  type EltT Acc a = Arrays a
  ifThenElse :: Exp Bool -> Acc a -> Acc a -> Acc a
ifThenElse = Exp Bool -> Acc a -> Acc a -> Acc a
forall a. Arrays a => Exp Bool -> Acc a -> Acc a -> Acc a
acond


-- | The 'match' operation is the core operation which enables embedded
-- pattern matching. It is applied to an n-ary scalar function, and
-- generates the necessary case-statements in the embedded code for each
-- argument. For example, given the function:
--
-- > example1 :: Exp (Maybe Bool) -> Exp Int
-- > example1 Nothing_ = 0
-- > example1 (Just_ False_) = 1
-- > example1 (Just_ True_) = 2
--
-- In order to use this function it must be applied to the 'match'
-- operator:
--
-- > match example1
--
-- Using the infix-flip operator ('Data.Function.&'), we can also write
-- case statements inline. For example, instead of this:
--
-- > example2 x = case f x of
-- >   Nothing_ -> ...      -- error: embedded pattern synonym...
-- >   Just_ y  -> ...      -- ...used outside of 'match' context
--
-- This can be written instead as:
--
-- > example3 x = f x & match \case
-- >   Nothing_ -> ...
-- >   Just_ y  -> ...
--
-- And utilising the @LambdaCase@ and @BlockArguments@ syntactic extensions.
--
-- The Template Haskell splice 'Data.Array.Accelerate.mkPattern' (or
-- 'Data.Array.Accelerate.mkPatterns') can be used to generate the pattern
-- synonyms for a given Haskell'98 sum or product data type. For example:
--
-- > data Option a = None | Some a
-- >   deriving (Generic, Elt)
-- >
-- > mkPattern ''Option
--
-- Which can then be used such as:
--
-- > isNone :: Elt a => Exp (Option a) -> Exp Bool
-- > isNone = match \case
-- >   None_   -> True_
-- >   Some_{} -> False_
--
-- @since 1.3.0.0
--
match :: Matching f => f -> f
match :: f -> f
match f
f = (Args f -> Exp (ResultT f)) -> (Args f -> Args f) -> f
forall a f.
Matching a =>
(Args f -> Exp (ResultT a)) -> (Args a -> Args f) -> a
mkFun (f -> Args f -> Exp (ResultT f)
forall a. Matching a => a -> Args a -> Exp (ResultT a)
mkMatch f
f) Args f -> Args f
forall a. a -> a
id

data Args f where
  (:->)  :: Exp a -> Args b -> Args (Exp a -> b)
  Result :: Args (Exp a)

class Matching a where
  type ResultT a
  mkMatch :: a -> Args a -> Exp (ResultT a)
  mkFun   :: (Args f -> Exp (ResultT a))
          -> (Args a -> Args f)
          -> a

instance Elt a => Matching (Exp a) where
  type ResultT (Exp a) = a

  mkFun :: (Args f -> Exp (ResultT (Exp a)))
-> (Args (Exp a) -> Args f) -> Exp a
mkFun Args f -> Exp (ResultT (Exp a))
f Args (Exp a) -> Args f
k = Args f -> Exp (ResultT (Exp a))
f (Args (Exp a) -> Args f
k Args (Exp a)
forall b. Args (Exp b)
Result)
  mkMatch :: Exp a -> Args (Exp a) -> Exp (ResultT (Exp a))
mkMatch (Exp SmartExp (EltR a)
e) Args (Exp a)
Result =
    case SmartExp (EltR a)
e of
      SmartExp (Match TagR (EltR a)
_ SmartExp (EltR a)
x) -> SmartExp (EltR a) -> Exp a
forall t. SmartExp (EltR t) -> Exp t
Exp SmartExp (EltR a)
x
      SmartExp (EltR a)
_                    -> SmartExp (EltR a) -> Exp a
forall t. SmartExp (EltR t) -> Exp t
Exp SmartExp (EltR a)
e

instance (Elt e, Matching r) => Matching (Exp e -> r) where
  type ResultT (Exp e -> r) = ResultT r

  mkFun :: (Args f -> Exp (ResultT (Exp e -> r)))
-> (Args (Exp e -> r) -> Args f) -> Exp e -> r
mkFun Args f -> Exp (ResultT (Exp e -> r))
f Args (Exp e -> r) -> Args f
k Exp e
x = (Args f -> Exp (ResultT r)) -> (Args r -> Args f) -> r
forall a f.
Matching a =>
(Args f -> Exp (ResultT a)) -> (Args a -> Args f) -> a
mkFun Args f -> Exp (ResultT r)
Args f -> Exp (ResultT (Exp e -> r))
f (\Args r
xs -> Args (Exp e -> r) -> Args f
k (Exp e
x Exp e -> Args r -> Args (Exp e -> r)
forall a b. Exp a -> Args b -> Args (Exp a -> b)
:-> Args r
xs))
  mkMatch :: (Exp e -> r) -> Args (Exp e -> r) -> Exp (ResultT (Exp e -> r))
mkMatch Exp e -> r
f (x :: Exp a
x@(Exp SmartExp (EltR a)
p) :-> Args b
xs) =
    case SmartExp (EltR a)
p of
      -- This first case is used when we have nested calls to 'match'
      SmartExp Match{} -> r -> Args r -> Exp (ResultT r)
forall a. Matching a => a -> Args a -> Exp (ResultT a)
mkMatch (Exp e -> r
f Exp e
Exp a
x) Args r
Args b
xs

      -- If there is only a single alternative, we can elide the case
      -- statement at this point. This can occur when pattern matching on
      -- product types.
      SmartExp (EltR a)
_ -> case [(TagR (EltR e), SmartExp (EltR (ResultT r)))]
rhs of
             [(TagR (EltR e)
_,SmartExp (EltR (ResultT r))
r)] -> SmartExp (EltR (ResultT r)) -> Exp (ResultT r)
forall t. SmartExp (EltR t) -> Exp t
Exp SmartExp (EltR (ResultT r))
r
             [(TagR (EltR e), SmartExp (EltR (ResultT r)))]
_       -> SmartExp (EltR (ResultT r)) -> Exp (ResultT r)
forall t. SmartExp (EltR t) -> Exp t
Exp (PreSmartExp SmartAcc SmartExp (EltR (ResultT r))
-> SmartExp (EltR (ResultT r))
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (SmartExp (EltR e)
-> [(TagR (EltR e), SmartExp (EltR (ResultT r)))]
-> PreSmartExp SmartAcc SmartExp (EltR (ResultT r))
forall (exp :: * -> *) a b (acc :: * -> *).
exp a -> [(TagR a, exp b)] -> PreSmartExp acc exp b
Case SmartExp (EltR e)
SmartExp (EltR a)
p [(TagR (EltR e), SmartExp (EltR (ResultT r)))]
rhs))
    where
      rhs :: [(TagR (EltR e), SmartExp (EltR (ResultT r)))]
rhs = [ (TagR (EltR e)
tag, Exp (ResultT r) -> SmartExp (EltR (ResultT r))
forall e. Exp e -> SmartExp (EltR e)
unExp (r -> Args r -> Exp (ResultT r)
forall a. Matching a => a -> Args a -> Exp (ResultT a)
mkMatch (Exp e -> r
f Exp e
x') Args r
Args b
xs))
            | TagR (EltR e)
tag <- Elt e => [TagR (EltR e)]
forall a. Elt a => [TagR (EltR a)]
tagsR @e
            , let x' :: Exp e
x' = SmartExp (EltR e) -> Exp e
forall t. SmartExp (EltR t) -> Exp t
Exp (PreSmartExp SmartAcc SmartExp (EltR e) -> SmartExp (EltR e)
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (TagR (EltR e)
-> SmartExp (EltR e) -> PreSmartExp SmartAcc SmartExp (EltR e)
forall t (exp :: * -> *) (acc :: * -> *).
TagR t -> exp t -> PreSmartExp acc exp t
Match TagR (EltR e)
tag SmartExp (EltR e)
SmartExp (EltR a)
p)) ]


-- Scalar iteration
-- ----------------

-- | Repeatedly apply a function a fixed number of times
--
iterate
    :: Elt a
    => Exp Int
    -> (Exp a -> Exp a)
    -> Exp a
    -> Exp a
iterate :: Exp Int -> (Exp a -> Exp a) -> Exp a -> Exp a
iterate Exp Int
n Exp a -> Exp a
f Exp a
z
  = let step :: Exp (Int, a) -> Exp (Int, a)
step (T2 Exp Int
i Exp a
acc) = Exp Int -> Exp a -> Exp (Int, a)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 (Exp Int
iExp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+Exp Int
1) (Exp a -> Exp a
f Exp a
acc)
     in Exp (Int, a) -> Exp a
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp b
snd (Exp (Int, a) -> Exp a) -> Exp (Int, a) -> Exp a
forall a b. (a -> b) -> a -> b
$ (Exp (Int, a) -> Exp Bool)
-> (Exp (Int, a) -> Exp (Int, a)) -> Exp (Int, a) -> Exp (Int, a)
forall e.
Elt e =>
(Exp e -> Exp Bool) -> (Exp e -> Exp e) -> Exp e -> Exp e
while (\Exp (Int, a)
v -> Exp (Int, a) -> Exp Int
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp a
fst Exp (Int, a)
v Exp Int -> Exp Int -> Exp Bool
forall a. Ord a => Exp a -> Exp a -> Exp Bool
< Exp Int
n) Exp (Int, a) -> Exp (Int, a)
step (Exp Int -> Exp a -> Exp (Int, a)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Exp Int
0 Exp a
z)


-- Scalar bulk operations
-- ----------------------

-- | Reduce along an innermost slice of an array /sequentially/, by applying a
-- binary operator to a starting value and the array from left to right.
--
sfoldl :: (Shape sh, Elt a, Elt b)
       => (Exp a -> Exp b -> Exp a)
       -> Exp a
       -> Exp sh
       -> Acc (Array (sh :. Int) b)
       -> Exp a
sfoldl :: (Exp a -> Exp b -> Exp a)
-> Exp a -> Exp sh -> Acc (Array (sh :. Int) b) -> Exp a
sfoldl Exp a -> Exp b -> Exp a
f Exp a
z Exp sh
ix Acc (Array (sh :. Int) b)
xs
  = let n :: Exp Int
n               = Exp (sh :. Int) -> Exp Int
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp a
indexHead (Acc (Array (sh :. Int) b) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) b)
xs)
        step :: Exp (Int, a) -> Exp (Int, a)
step (T2 Exp Int
i Exp a
acc) = Exp Int -> Exp a -> Exp (Int, a)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 (Exp Int
iExp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+Exp Int
1) (Exp a
acc Exp a -> Exp b -> Exp a
`f` (Acc (Array (sh :. Int) b)
xs Acc (Array (sh :. Int) b) -> Exp (sh :. Int) -> Exp b
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! (Exp sh
ix Exp sh -> Exp Int -> Exp (sh :. Int)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp Int
i)))
     in Exp (Int, a) -> Exp a
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp b
snd (Exp (Int, a) -> Exp a) -> Exp (Int, a) -> Exp a
forall a b. (a -> b) -> a -> b
$ (Exp (Int, a) -> Exp Bool)
-> (Exp (Int, a) -> Exp (Int, a)) -> Exp (Int, a) -> Exp (Int, a)
forall e.
Elt e =>
(Exp e -> Exp Bool) -> (Exp e -> Exp e) -> Exp e -> Exp e
while (\Exp (Int, a)
v -> Exp (Int, a) -> Exp Int
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp a
fst Exp (Int, a)
v Exp Int -> Exp Int -> Exp Bool
forall a. Ord a => Exp a -> Exp a -> Exp Bool
< Exp Int
n) Exp (Int, a) -> Exp (Int, a)
step (Exp Int -> Exp a -> Exp (Int, a)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Exp Int
0 Exp a
z)


-- Tuples
-- ------

-- |Extract the first component of a scalar pair.
--
fst :: (Elt a, Elt b) => Exp (a, b) -> Exp a
fst :: Exp (a, b) -> Exp a
fst (T2 Exp a
a Exp b
_) = Exp a
a

-- |Extract the first component of an array pair.
{-# NOINLINE[1] afst #-}
afst :: (Arrays a, Arrays b) => Acc (a, b) -> Acc a
afst :: Acc (a, b) -> Acc a
afst (T2 Acc a
a Acc b
_) = Acc a
a

-- |Extract the second component of a scalar pair.
--
snd :: (Elt a, Elt b) => Exp (a, b) -> Exp b
snd :: Exp (a, b) -> Exp b
snd (T2 Exp a
_ Exp b
b) = Exp b
b

-- | Extract the second component of an array pair
asnd :: (Arrays a, Arrays b) => Acc (a, b) -> Acc b
asnd :: Acc (a, b) -> Acc b
asnd (T2 Acc a
_ Acc b
b) = Acc b
b

-- |Converts an uncurried function to a curried function.
--
curry :: Lift f (f a, f b) => (f (Plain (f a), Plain (f b)) -> f c) -> f a -> f b -> f c
curry :: (f (Plain (f a), Plain (f b)) -> f c) -> f a -> f b -> f c
curry f (Plain (f a), Plain (f b)) -> f c
f f a
x f b
y = f (Plain (f a), Plain (f b)) -> f c
f ((f a, f b) -> f (Plain (f a, f b))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (f a
x, f b
y))

-- |Converts a curried function to a function on pairs.
--
uncurry :: Unlift f (f a, f b) => (f a -> f b -> f c) -> f (Plain (f a), Plain (f b)) -> f c
uncurry :: (f a -> f b -> f c) -> f (Plain (f a), Plain (f b)) -> f c
uncurry f a -> f b -> f c
f f (Plain (f a), Plain (f b))
t = let (f a
x, f b
y) = f (Plain (f a, f b)) -> (f a, f b)
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift f (Plain (f a), Plain (f b))
f (Plain (f a, f b))
t in f a -> f b -> f c
f f a
x f b
y


-- Shapes and indices
-- ------------------

-- | The one index for a rank-0 array.
--
index0 :: Exp Z
index0 :: Exp DIM0
index0 = DIM0 -> Exp (Plain DIM0)
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift DIM0
Z

-- | Turn an 'Int' expression into a rank-1 indexing expression.
--
index1 :: Elt i => Exp i -> Exp (Z :. i)
index1 :: Exp i -> Exp (DIM0 :. i)
index1 Exp i
i = (DIM0 :. Exp i) -> Exp (Plain (DIM0 :. Exp i))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (DIM0
Z DIM0 -> Exp i -> DIM0 :. Exp i
forall tail head. tail -> head -> tail :. head
:. Exp i
i)

-- | Turn a rank-1 indexing expression into an 'Int' expression.
--
unindex1 :: Elt i => Exp (Z :. i) -> Exp i
unindex1 :: Exp (DIM0 :. i) -> Exp i
unindex1 Exp (DIM0 :. i)
ix = let DIM0
Z :. Exp i
i = Exp (Plain (DIM0 :. Exp i)) -> DIM0 :. Exp i
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp (DIM0 :. i)
Exp (Plain (DIM0 :. Exp i))
ix in Exp i
i

-- | Creates a rank-2 index from two Exp Int`s
--
index2
    :: Elt i
    => Exp i
    -> Exp i
    -> Exp (Z :. i :. i)
index2 :: Exp i -> Exp i -> Exp ((DIM0 :. i) :. i)
index2 Exp i
i Exp i
j = ((DIM0 :. Exp i) :. Exp i)
-> Exp (Plain ((DIM0 :. Exp i) :. Exp i))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (DIM0
Z DIM0 -> Exp i -> DIM0 :. Exp i
forall tail head. tail -> head -> tail :. head
:. Exp i
i (DIM0 :. Exp i) -> Exp i -> (DIM0 :. Exp i) :. Exp i
forall tail head. tail -> head -> tail :. head
:. Exp i
j)

-- | Destructs a rank-2 index to an Exp tuple of two Int`s.
--
unindex2
    :: Elt i
    => Exp (Z :. i :. i)
    -> Exp (i, i)
unindex2 :: Exp ((DIM0 :. i) :. i) -> Exp (i, i)
unindex2 (Exp DIM0
Z_ ::. Exp i
i ::. Exp i
j) = Exp i -> Exp i -> Exp (i, i)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Exp i
i Exp i
j

-- | Create a rank-3 index from three Exp Int`s
--
index3
    :: Elt i
    => Exp i
    -> Exp i
    -> Exp i
    -> Exp (Z :. i :. i :. i)
index3 :: Exp i -> Exp i -> Exp i -> Exp (((DIM0 :. i) :. i) :. i)
index3 Exp i
k Exp i
j Exp i
i = Exp DIM0
Z_ Exp DIM0 -> Exp i -> Exp (DIM0 :. i)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp i
k Exp (DIM0 :. i) -> Exp i -> Exp ((DIM0 :. i) :. i)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp i
j Exp ((DIM0 :. i) :. i) -> Exp i -> Exp (((DIM0 :. i) :. i) :. i)
forall a b. (Elt a, Elt b) => Exp a -> Exp b -> Exp (a :. b)
::. Exp i
i

-- | Destruct a rank-3 index into an Exp tuple of Int`s
unindex3
    :: Elt i
    => Exp (Z :. i :. i :. i)
    -> Exp (i, i, i)
unindex3 :: Exp (((DIM0 :. i) :. i) :. i) -> Exp (i, i, i)
unindex3 (Exp DIM0
Z_ ::. Exp i
k ::. Exp i
j ::. Exp i
i) = Exp i -> Exp i -> Exp i -> Exp (i, i, i)
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con x0 -> con x1 -> con x2 -> con (x0, x1, x2)
T3 Exp i
k Exp i
j Exp i
i


-- Array operations with a scalar result
-- -------------------------------------

-- | Extract the element of a singleton array.
--
-- > the xs  ==  xs ! Z
--
the :: Elt e => Acc (Scalar e) -> Exp e
the :: Acc (Scalar e) -> Exp e
the = (Acc (Scalar e) -> Exp DIM0 -> Exp e
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
!Exp DIM0
index0)

-- | Test whether an array is empty.
--
null :: (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool
null :: Acc (Array sh e) -> Exp Bool
null Acc (Array sh e)
arr = Acc (Array sh e) -> Exp Int
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
size Acc (Array sh e)
arr Exp Int -> Exp Int -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp Int
0

-- | Get the length of a vector.
--
length :: Elt e => Acc (Vector e) -> Exp Int
length :: Acc (Vector e) -> Exp Int
length = Exp (DIM0 :. Int) -> Exp Int
forall i. Elt i => Exp (DIM0 :. i) -> Exp i
unindex1 (Exp (DIM0 :. Int) -> Exp Int)
-> (Acc (Vector e) -> Exp (DIM0 :. Int))
-> Acc (Vector e)
-> Exp Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acc (Vector e) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape


-- Operations to facilitate irregular data parallelism
-- ---------------------------------------------------

-- | A recipe for generating flattened implementations of some kinds of
-- irregular nested parallelism. Given two functions that:
--
--   (1) for each source element, determine how many target
--       elements it expands into; and
--
--   (2) computes a particular target element based on a source element and
--       the target element index associated with the source
--
-- The following example implements the Sieve of Eratosthenes,
-- a contraction style algorithm which first computes all primes less than
-- /sqrt n/, then uses this intermediate result to sieve away all numbers
-- in the range /[sqrt n .. n]/. The 'expand' function is used to calculate
-- and flatten the sieves. For each prime /p/ and upper limit /c2/,
-- function /sz/ computes the number of contributions in the sieve. Then,
-- for each prime /p/ and sieve index /i/, the function /get/ computes the
-- sieve contribution. The final step produces all the new primes in the
-- interval /[c1 .. c2]/.
--
-- >>> :{
--   primes :: Exp Int -> Acc (Vector Int)
--   primes n = afst loop
--     where
--       c0    = unit 2
--       a0    = use $ fromList (Z:.0) []
--       limit = truncate (sqrt (fromIntegral (n+1) :: Exp Float))
--       loop  = awhile
--                 (\(T2 _   c) -> map (< n+1) c)
--                 (\(T2 old c) ->
--                   let c1 = the c
--                       c2 = c1 < limit ? ( c1*c1, n+1 )
--                       --
--                       sieves =
--                         let sz p    = (c2 - p) `quot` p
--                             get p i = (2+i)*p
--                         in
--                         map (subtract c1) (expand sz get old)
--                       --
--                       new =
--                         let m     = c2-c1
--                             put i = let s = sieves ! i
--                                      in s >= 0 && s < m ? (Just_ (I1 s), Nothing_)
--                         in
--                         afst
--                           $ filter (> 0)
--                           $ permute const (enumFromN (I1 m) c1) put
--                           $ fill (shape sieves) 0
--                    in
--                    T2 (old ++ new) (unit c2))
--                 (T2 a0 c0)
-- :}
--
-- >>> run $ primes 100
-- Vector (Z :. 25) [2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71,73,79,83,89,97]
--
-- Inspired by the paper /Data-Parallel Flattening by Expansion/ by Martin
-- Elsman, Troels Henriksen, and Niels Gustav Westphal Serup, ARRAY'19.
--
-- @since 1.3.0.0
--
expand :: (Elt a, Elt b)
       => (Exp a -> Exp Int)
       -> (Exp a -> Exp Int -> Exp b)
       -> Acc (Vector a)
       -> Acc (Vector b)
expand :: (Exp a -> Exp Int)
-> (Exp a -> Exp Int -> Exp b) -> Acc (Vector a) -> Acc (Vector b)
expand Exp a -> Exp Int
f Exp a -> Exp Int -> Exp b
g Acc (Vector a)
xs =
  let
      szs :: Acc (Array (DIM0 :. Int) Int)
szs           = (Exp a -> Exp Int)
-> Acc (Vector a) -> Acc (Array (DIM0 :. Int) Int)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp a -> Exp Int
f Acc (Vector a)
xs
      T2 Acc (Array (DIM0 :. Int) Int)
offset Acc (Scalar Int)
len = (Exp Int -> Exp Int -> Exp Int)
-> Exp Int
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) Int, Scalar Int)
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanl' Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
(+) Exp Int
0 Acc (Array (DIM0 :. Int) Int)
szs
      m :: Exp Int
m             = Acc (Scalar Int) -> Exp Int
forall e. Elt e => Acc (Scalar e) -> Exp e
the Acc (Scalar Int)
len
  in
  if Acc (Vector a) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
length Acc (Vector a)
xs Exp Int -> Exp Int -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp Int
0 Exp Bool -> Exp Bool -> Exp Bool
|| Exp Int
m Exp Int -> Exp Int -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp Int
0
     then Vector b -> Acc (Vector b)
forall arrays. Arrays arrays => arrays -> Acc arrays
use (Vector b -> Acc (Vector b)) -> Vector b -> Acc (Vector b)
forall a b. (a -> b) -> a -> b
$ (DIM0 :. Int) -> [b] -> Vector b
forall sh e. (Shape sh, Elt e) => sh -> [e] -> Array sh e
fromList (DIM0
ZDIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:.Int
0) []
     else
      let
          n :: Exp Int
n          = Exp Int
m Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
+ Exp Int
1
          put :: Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
put Exp (DIM0 :. Int)
ix     = Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 (Acc (Array (DIM0 :. Int) Int)
offset Acc (Array (DIM0 :. Int) Int) -> Exp (DIM0 :. Int) -> Exp Int
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh e) -> Exp sh -> Exp e
! Exp (DIM0 :. Int)
ix))

          head_flags :: Acc (Vector Int)
          head_flags :: Acc (Array (DIM0 :. Int) Int)
head_flags = (Exp Int -> Exp Int -> Exp Int)
-> Acc (Array (DIM0 :. Int) Int)
-> (Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) Int)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp Int -> Exp Int -> Exp Int
forall a b. a -> b -> a
const (Exp (DIM0 :. Int) -> Exp Int -> Acc (Array (DIM0 :. Int) Int)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
n) Exp Int
0) Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
put (Exp (DIM0 :. Int) -> Exp Int -> Acc (Array (DIM0 :. Int) Int)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Acc (Array (DIM0 :. Int) Int) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (DIM0 :. Int) Int)
szs) Exp Int
1)

          idxs :: Acc (Array (DIM0 :. Int) Int)
idxs       = (Exp Int -> Exp Int)
-> Acc (Array (DIM0 :. Int) Int) -> Acc (Array (DIM0 :. Int) Int)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map (Exp Int -> Exp Int -> Exp Int
forall a. Num a => Exp a -> Exp a -> Exp a
subtract Exp Int
1)
                     (Acc (Array (DIM0 :. Int) Int) -> Acc (Array (DIM0 :. Int) Int))
-> Acc (Array (DIM0 :. Int) Int) -> Acc (Array (DIM0 :. Int) Int)
forall a b. (a -> b) -> a -> b
$ (Exp (Int, Int) -> Exp Int)
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) Int)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (Int, Int) -> Exp Int
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp b
snd
                     (Acc (Array (DIM0 :. Int) (Int, Int))
 -> Acc (Array (DIM0 :. Int) Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) Int)
forall a b. (a -> b) -> a -> b
$ (Exp (Int, Int) -> Exp (Int, Int) -> Exp (Int, Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 ((Exp Int -> Exp Int -> Exp Int)
-> Exp (Int, Int) -> Exp (Int, Int) -> Exp (Int, Int)
forall e i.
(Elt e, Num i, Bits i) =>
(Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
segmentedL Exp Int -> Exp Int -> Exp Int
forall a. Num a => a -> a -> a
(+))
                     (Acc (Array (DIM0 :. Int) (Int, Int))
 -> Acc (Array (DIM0 :. Int) (Int, Int)))
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall a b. (a -> b) -> a -> b
$ Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip Acc (Array (DIM0 :. Int) Int)
head_flags
                     (Acc (Array (DIM0 :. Int) Int)
 -> Acc (Array (DIM0 :. Int) (Int, Int)))
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall a b. (a -> b) -> a -> b
$ Exp (DIM0 :. Int) -> Exp Int -> Acc (Array (DIM0 :. Int) Int)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
m) Exp Int
1

          iotas :: Acc (Array (DIM0 :. Int) Int)
iotas      = (Exp (Int, Int) -> Exp Int)
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) Int)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (Int, Int) -> Exp Int
forall a b. (Elt a, Elt b) => Exp (a, b) -> Exp b
snd
                     (Acc (Array (DIM0 :. Int) (Int, Int))
 -> Acc (Array (DIM0 :. Int) Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) Int)
forall a b. (a -> b) -> a -> b
$ (Exp (Int, Int) -> Exp (Int, Int) -> Exp (Int, Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall sh a.
(Shape sh, Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 ((Exp Int -> Exp Int -> Exp Int)
-> Exp (Int, Int) -> Exp (Int, Int) -> Exp (Int, Int)
forall e i.
(Elt e, Num i, Bits i) =>
(Exp e -> Exp e -> Exp e) -> Exp (i, e) -> Exp (i, e) -> Exp (i, e)
segmentedL Exp Int -> Exp Int -> Exp Int
forall a b. a -> b -> a
const)
                     (Acc (Array (DIM0 :. Int) (Int, Int))
 -> Acc (Array (DIM0 :. Int) (Int, Int)))
-> Acc (Array (DIM0 :. Int) (Int, Int))
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall a b. (a -> b) -> a -> b
$ Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh (a, b))
zip Acc (Array (DIM0 :. Int) Int)
head_flags
                     (Acc (Array (DIM0 :. Int) Int)
 -> Acc (Array (DIM0 :. Int) (Int, Int)))
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) (Int, Int))
forall a b. (a -> b) -> a -> b
$ (Exp Int -> Exp Int -> Exp Int)
-> Acc (Array (DIM0 :. Int) Int)
-> (Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int)))
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Array (DIM0 :. Int) Int)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
(Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute Exp Int -> Exp Int -> Exp Int
forall a b. a -> b -> a
const (Exp (DIM0 :. Int) -> Exp Int -> Acc (Array (DIM0 :. Int) Int)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (Exp Int -> Exp (DIM0 :. Int)
forall x0. Elt x0 => Exp x0 -> Exp (DIM0 :. x0)
I1 Exp Int
n) Exp Int
forall e. Elt e => Exp e
undef) Exp (DIM0 :. Int) -> Exp (Maybe (DIM0 :. Int))
put
                     (Acc (Array (DIM0 :. Int) Int) -> Acc (Array (DIM0 :. Int) Int))
-> Acc (Array (DIM0 :. Int) Int) -> Acc (Array (DIM0 :. Int) Int)
forall a b. (a -> b) -> a -> b
$ Exp (DIM0 :. Int) -> Exp Int -> Acc (Array (DIM0 :. Int) Int)
forall sh e.
(Shape sh, Num e, FromIntegral Int e) =>
Exp sh -> Exp e -> Acc (Array sh e)
enumFromN (Acc (Vector a) -> Exp (DIM0 :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Vector a)
xs) Exp Int
0
      in
      (Exp a -> Exp Int -> Exp b)
-> Acc (Vector a)
-> Acc (Array (DIM0 :. Int) Int)
-> Acc (Vector b)
forall sh a b c.
(Shape sh, Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
zipWith Exp a -> Exp Int -> Exp b
g (Acc (Array (DIM0 :. Int) Int) -> Acc (Vector a) -> Acc (Vector a)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array sh Int) -> Acc (Vector e) -> Acc (Array sh e)
gather Acc (Array (DIM0 :. Int) Int)
iotas Acc (Vector a)
xs) Acc (Array (DIM0 :. Int) Int)
idxs


{--
-- Sequence operations
-- -------------------

-- | Reduce a sequence by appending all the shapes and all the elements in two
-- separate vectors.
--
fromSeq :: (Shape sh, Elt a) => Seq [Array sh a] -> Seq (Vector sh, Vector a)
fromSeq = foldSeqFlatten f (lift (emptyArray, emptyArray))
  where
    f x sh1 a1 =
      let (sh0, a0) = unlift x
      in lift (sh0 ++ sh1, a0 ++ a1)


fromSeqElems :: (Shape sh, Elt a) => Seq [Array sh a] -> Seq (Vector a)
fromSeqElems = foldSeqFlatten f emptyArray
  where
    f a0 _ a1 = a0 ++ a1

fromSeqShapes :: (Shape sh, Elt a) => Seq [Array sh a] -> Seq (Vector sh)
fromSeqShapes = foldSeqFlatten f emptyArray
  where
    f sh0 sh1 _ = sh0 ++ sh1

-- | Sequence an array on the innermost dimension.
--
toSeqInner :: (Shape sh, Elt a) => Acc (Array (sh :. Int) a) -> Seq [Array sh a]
toSeqInner a = toSeq (Any :. Split) a

-- | Sequence a 2-dimensional array on the outermost dimension.
--
toSeqOuter2 :: Elt a => Acc (Array DIM2 a) -> Seq [Array DIM1 a]
toSeqOuter2 a = toSeq (Z :. Split :. All) a

-- | Sequence a 3-dimensional array on the outermost dimension.
toSeqOuter3 :: Elt a => Acc (Array DIM3 a) -> Seq [Array DIM2 a]
toSeqOuter3 a = toSeq (Z :. Split :. All :. All) a

-- | Generate a scalar sequence of a fixed given length, by applying
-- the given scalar function at each index.
generateSeq :: Elt a => Exp Int -> (Exp Int -> Exp a) -> Seq [Scalar a]
generateSeq n f = toSeq (Z :. Split) (generate (index1 n) (f . unindex1))
--}

-- Utilities
-- ---------

emptyArray :: (Shape sh, Elt e) => Acc (Array sh e)
emptyArray :: Acc (Array sh e)
emptyArray = Exp sh -> Exp e -> Acc (Array sh e)
forall sh e.
(Shape sh, Elt e) =>
Exp sh -> Exp e -> Acc (Array sh e)
fill (sh -> Exp sh
forall e. (HasCallStack, Elt e) => e -> Exp e
constant sh
forall sh. Shape sh => sh
empty) Exp e
forall e. Elt e => Exp e
undef


-- Lenses
-- ------
--
-- Imported from `lens-accelerate` (which provides more general Field instances)
--
_1 :: forall sh. Elt sh => Lens' (Exp (sh:.Int)) (Exp Int)
_1 :: Lens' (Exp (sh :. Int)) (Exp Int)
_1 = (Exp (sh :. Int) -> Exp Int)
-> (Exp (sh :. Int) -> Exp Int -> Exp (sh :. Int))
-> Lens' (Exp (sh :. Int)) (Exp Int)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (\Exp (sh :. Int)
ix   -> let Exp sh
_  :. Exp Int
x = Exp (Plain (Exp sh :. Exp Int)) -> Exp sh :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp (sh :. Int)
Exp (Plain (Exp sh :. Exp Int))
ix :: Exp sh :. Exp Int in Exp Int
x)
          (\Exp (sh :. Int)
ix Exp Int
x -> let Exp sh
sh :. Exp Int
_ = Exp (Plain (Exp sh :. Exp Int)) -> Exp sh :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp (sh :. Int)
Exp (Plain (Exp sh :. Exp Int))
ix :: Exp sh :. Exp Int in (Exp sh :. Exp Int) -> Exp (Plain (Exp sh :. Exp Int))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (Exp sh
sh Exp sh -> Exp Int -> Exp sh :. Exp Int
forall tail head. tail -> head -> tail :. head
:. Exp Int
x))

_2 :: forall sh. Elt sh => Lens' (Exp (sh:.Int:.Int)) (Exp Int)
_2 :: Lens' (Exp ((sh :. Int) :. Int)) (Exp Int)
_2 = (Exp ((sh :. Int) :. Int) -> Exp Int)
-> (Exp ((sh :. Int) :. Int)
    -> Exp Int -> Exp ((sh :. Int) :. Int))
-> Lens' (Exp ((sh :. Int) :. Int)) (Exp Int)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (\Exp ((sh :. Int) :. Int)
ix   -> let Exp sh
_  :. Exp Int
y :. Exp Int
_ = Exp (Plain ((Exp sh :. Exp Int) :. Exp Int))
-> (Exp sh :. Exp Int) :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp ((sh :. Int) :. Int)
Exp (Plain ((Exp sh :. Exp Int) :. Exp Int))
ix :: Exp sh :. Exp Int :. Exp Int in Exp Int
y)
          (\Exp ((sh :. Int) :. Int)
ix Exp Int
y -> let Exp sh
sh :. Exp Int
_ :. Exp Int
x = Exp (Plain ((Exp sh :. Exp Int) :. Exp Int))
-> (Exp sh :. Exp Int) :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp ((sh :. Int) :. Int)
Exp (Plain ((Exp sh :. Exp Int) :. Exp Int))
ix :: Exp sh :. Exp Int :. Exp Int in ((Exp sh :. Exp Int) :. Exp Int)
-> Exp (Plain ((Exp sh :. Exp Int) :. Exp Int))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (Exp sh
sh Exp sh -> Exp Int -> Exp sh :. Exp Int
forall tail head. tail -> head -> tail :. head
:. Exp Int
y (Exp sh :. Exp Int) -> Exp Int -> (Exp sh :. Exp Int) :. Exp Int
forall tail head. tail -> head -> tail :. head
:. Exp Int
x))

_3 :: forall sh. Elt sh => Lens' (Exp (sh:.Int:.Int:.Int)) (Exp Int)
_3 :: Lens' (Exp (((sh :. Int) :. Int) :. Int)) (Exp Int)
_3 = (Exp (((sh :. Int) :. Int) :. Int) -> Exp Int)
-> (Exp (((sh :. Int) :. Int) :. Int)
    -> Exp Int -> Exp (((sh :. Int) :. Int) :. Int))
-> Lens' (Exp (((sh :. Int) :. Int) :. Int)) (Exp Int)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (\Exp (((sh :. Int) :. Int) :. Int)
ix   -> let Exp sh
_  :. Exp Int
z :. Exp Int
_ :. Exp Int
_ = Exp (Plain (((Exp sh :. Exp Int) :. Exp Int) :. Exp Int))
-> ((Exp sh :. Exp Int) :. Exp Int) :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp (((sh :. Int) :. Int) :. Int)
Exp (Plain (((Exp sh :. Exp Int) :. Exp Int) :. Exp Int))
ix :: Exp sh :. Exp Int :. Exp Int :. Exp Int in Exp Int
z)
          (\Exp (((sh :. Int) :. Int) :. Int)
ix Exp Int
z -> let Exp sh
sh :. Exp Int
_ :. Exp Int
y :. Exp Int
x = Exp (Plain (((Exp sh :. Exp Int) :. Exp Int) :. Exp Int))
-> ((Exp sh :. Exp Int) :. Exp Int) :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp (((sh :. Int) :. Int) :. Int)
Exp (Plain (((Exp sh :. Exp Int) :. Exp Int) :. Exp Int))
ix :: Exp sh :. Exp Int :. Exp Int :. Exp Int in (((Exp sh :. Exp Int) :. Exp Int) :. Exp Int)
-> Exp (Plain (((Exp sh :. Exp Int) :. Exp Int) :. Exp Int))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (Exp sh
sh Exp sh -> Exp Int -> Exp sh :. Exp Int
forall tail head. tail -> head -> tail :. head
:. Exp Int
z (Exp sh :. Exp Int) -> Exp Int -> (Exp sh :. Exp Int) :. Exp Int
forall tail head. tail -> head -> tail :. head
:. Exp Int
y ((Exp sh :. Exp Int) :. Exp Int)
-> Exp Int -> ((Exp sh :. Exp Int) :. Exp Int) :. Exp Int
forall tail head. tail -> head -> tail :. head
:. Exp Int
x))