module Futhark.Representation.AST.Attributes.Reshape
       (
         -- * Basic tools
         newDim
       , newDims
       , newShape

         -- * Construction
       , shapeCoerce
       , repeatShapes

         -- * Execution
       , reshapeOuter
       , reshapeInner
       , repeatDims

         -- * Inspection
       , shapeCoercion

         -- * Simplification
       , fuseReshape
       , informReshape

         -- * Shape calculations
       , reshapeIndex
       , flattenIndex
       , unflattenIndex
       , sliceSizes
       )
       where

import Data.Foldable

import Prelude hiding (sum, product, quot)

import Futhark.Representation.AST.Attributes.Types
import Futhark.Representation.AST.Syntax
import Futhark.Util.IntegralExp

-- | The new dimension.
newDim :: DimChange d -> d
newDim :: DimChange d -> d
newDim (DimCoercion d
se) = d
se
newDim (DimNew      d
se) = d
se

-- | The new dimensions resulting from a reshape operation.
newDims :: ShapeChange d -> [d]
newDims :: ShapeChange d -> [d]
newDims = (DimChange d -> d) -> ShapeChange d -> [d]
forall a b. (a -> b) -> [a] -> [b]
map DimChange d -> d
forall d. DimChange d -> d
newDim

-- | The new shape resulting from a reshape operation.
newShape :: ShapeChange SubExp -> Shape
newShape :: ShapeChange SubExp -> Shape
newShape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape)
-> (ShapeChange SubExp -> [SubExp]) -> ShapeChange SubExp -> Shape
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeChange SubExp -> [SubExp]
forall d. ShapeChange d -> [d]
newDims

-- ^ Construct a 'Reshape' where all dimension changes are
-- 'DimCoercion's.
shapeCoerce :: [SubExp] -> VName -> Exp lore
shapeCoerce :: [SubExp] -> VName -> Exp lore
shapeCoerce [SubExp]
newdims VName
arr =
  BasicOp lore -> Exp lore
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp lore -> Exp lore) -> BasicOp lore -> Exp lore
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp lore
forall lore. ShapeChange SubExp -> VName -> BasicOp lore
Reshape ((SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion [SubExp]
newdims) VName
arr

-- | Construct a pair suitable for a 'Repeat'.
repeatShapes :: [Shape] -> Type -> ([Shape], Shape)
repeatShapes :: [Shape] -> Type -> ([Shape], Shape)
repeatShapes [Shape]
shapes Type
t =
  case Int -> [Shape] -> ([Shape], [Shape])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
t_rank [Shape]
shapes of
    ([Shape]
outer_shapes, [Shape
inner_shape]) ->
      ([Shape]
outer_shapes, Shape
inner_shape)
    ([Shape], [Shape])
_ ->
      ([Shape]
shapes [Shape] -> [Shape] -> [Shape]
forall a. [a] -> [a] -> [a]
++ Int -> Shape -> [Shape]
forall a. Int -> a -> [a]
replicate ([Shape] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
shapes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
t_rank) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape []), [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [])
  where t_rank :: Int
t_rank = Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
arrayRank Type
t

-- | Modify the shape of an array type as 'Repeat' would do
repeatDims :: [Shape] -> Shape -> Type -> Type
repeatDims :: [Shape] -> Shape -> Type -> Type
repeatDims [Shape]
shape Shape
innershape = (Shape -> Shape) -> Type -> Type
forall newshape oldshape u.
ArrayShape newshape =>
(oldshape -> newshape)
-> TypeBase oldshape u -> TypeBase newshape u
modifyArrayShape Shape -> Shape
repeatDims'
  where repeatDims' :: Shape -> Shape
repeatDims' (Shape [SubExp]
ds) =
          [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (([SubExp] -> [SubExp] -> [SubExp])
-> [[SubExp]] -> [[SubExp]] -> [[SubExp]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
(++) ((Shape -> [SubExp]) -> [Shape] -> [[SubExp]]
forall a b. (a -> b) -> [a] -> [b]
map Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims [Shape]
shape) ((SubExp -> [SubExp]) -> [SubExp] -> [[SubExp]]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ds)) [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++
          Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
innershape

-- | @reshapeOuter newshape n oldshape@ returns a 'Reshape' expression
-- that replaces the outer @n@ dimensions of @oldshape@ with @newshape@.
reshapeOuter :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter ShapeChange SubExp
newshape Int
n Shape
oldshape =
  ShapeChange SubExp
newshape ShapeChange SubExp -> ShapeChange SubExp -> ShapeChange SubExp
forall a. [a] -> [a] -> [a]
++ (SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
coercion_or_new (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
n (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
oldshape))
  where coercion_or_new :: d -> DimChange d
coercion_or_new
          | ShapeChange SubExp -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShapeChange SubExp
newshape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = d -> DimChange d
forall d. d -> DimChange d
DimCoercion
          | Bool
otherwise            = d -> DimChange d
forall d. d -> DimChange d
DimNew

-- | @reshapeInner newshape n oldshape@ returns a 'Reshape' expression
-- that replaces the inner @m-n@ dimensions (where @m@ is the rank of
-- @oldshape@) of @src@ with @newshape@.
reshapeInner :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeInner :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeInner ShapeChange SubExp
newshape Int
n Shape
oldshape =
  (SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
coercion_or_new (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
n (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
oldshape)) ShapeChange SubExp -> ShapeChange SubExp -> ShapeChange SubExp
forall a. [a] -> [a] -> [a]
++ ShapeChange SubExp
newshape
  where coercion_or_new :: d -> DimChange d
coercion_or_new
          | ShapeChange SubExp -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShapeChange SubExp
newshape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
n = d -> DimChange d
forall d. d -> DimChange d
DimCoercion
          | Bool
otherwise              = d -> DimChange d
forall d. d -> DimChange d
DimNew
        m :: Int
m = Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
oldshape

-- | If the shape change is nothing but shape coercions, return the new dimensions.  Otherwise, return
-- 'Nothing'.
shapeCoercion :: ShapeChange d -> Maybe [d]
shapeCoercion :: ShapeChange d -> Maybe [d]
shapeCoercion = (DimChange d -> Maybe d) -> ShapeChange d -> Maybe [d]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimChange d -> Maybe d
forall a. DimChange a -> Maybe a
dimCoercion
  where dimCoercion :: DimChange a -> Maybe a
dimCoercion (DimCoercion a
d) = a -> Maybe a
forall a. a -> Maybe a
Just a
d
        dimCoercion (DimNew      a
_) = Maybe a
forall a. Maybe a
Nothing

-- | @fuseReshape s1 s2@ creates a new 'ShapeChange' that is
-- semantically the same as first applying @s1@ and then @s2@.  This
-- may take advantage of properties of 'DimCoercion' versus 'DimNew'
-- to preserve information.
fuseReshape :: Eq d => ShapeChange d -> ShapeChange d -> ShapeChange d
fuseReshape :: ShapeChange d -> ShapeChange d -> ShapeChange d
fuseReshape ShapeChange d
s1 ShapeChange d
s2
  | ShapeChange d -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShapeChange d
s1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ShapeChange d -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShapeChange d
s2 =
      (DimChange d -> DimChange d -> DimChange d)
-> ShapeChange d -> ShapeChange d -> ShapeChange d
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith DimChange d -> DimChange d -> DimChange d
forall d. Eq d => DimChange d -> DimChange d -> DimChange d
comb ShapeChange d
s1 ShapeChange d
s2
  where comb :: DimChange d -> DimChange d -> DimChange d
comb (DimNew d
_)       (DimCoercion d
d2) =
          d -> DimChange d
forall d. d -> DimChange d
DimNew d
d2
        comb (DimCoercion d
d1) (DimNew d
d2)
          | d
d1 d -> d -> Bool
forall a. Eq a => a -> a -> Bool
== d
d2  = d -> DimChange d
forall d. d -> DimChange d
DimCoercion d
d2
          | Bool
otherwise = d -> DimChange d
forall d. d -> DimChange d
DimNew d
d2
        comb DimChange d
_                DimChange d
d2 =
          DimChange d
d2
-- TODO: intelligently handle case where s1 is a prefix of s2.
fuseReshape ShapeChange d
_ ShapeChange d
s2 = ShapeChange d
s2

-- | Given concrete information about the shape of the source array,
-- convert some 'DimNew's into 'DimCoercion's.
informReshape :: Eq d => [d] -> ShapeChange d -> ShapeChange d
informReshape :: [d] -> ShapeChange d -> ShapeChange d
informReshape [d]
shape ShapeChange d
sc
  | [d] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [d]
shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ShapeChange d -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShapeChange d
sc =
    (d -> DimChange d -> DimChange d)
-> [d] -> ShapeChange d -> ShapeChange d
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith d -> DimChange d -> DimChange d
forall d. Eq d => d -> DimChange d -> DimChange d
inform [d]
shape ShapeChange d
sc
  where inform :: d -> DimChange d -> DimChange d
inform d
d1 (DimNew d
d2)
          | d
d1 d -> d -> Bool
forall a. Eq a => a -> a -> Bool
== d
d2  = d -> DimChange d
forall d. d -> DimChange d
DimCoercion d
d2
        inform d
_ DimChange d
dc =
          DimChange d
dc
informReshape [d]
_ ShapeChange d
sc = ShapeChange d
sc

-- | @reshapeIndex to_dims from_dims is@ transforms the index list
-- @is@ (which is into an array of shape @from_dims@) into an index
-- list @is'@, which is into an array of shape @to_dims@.  @is@ must
-- have the same length as @from_dims@, and @is'@ will have the same
-- length as @to_dims@.
reshapeIndex :: IntegralExp num =>
                [num] -> [num] -> [num] -> [num]
reshapeIndex :: [num] -> [num] -> [num] -> [num]
reshapeIndex [num]
to_dims [num]
from_dims [num]
is =
  [num] -> num -> [num]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex [num]
to_dims (num -> [num]) -> num -> [num]
forall a b. (a -> b) -> a -> b
$ [num] -> [num] -> num
forall num. IntegralExp num => [num] -> [num] -> num
flattenIndex [num]
from_dims [num]
is

-- | @unflattenIndex dims i@ computes a list of indices into an array
-- with dimension @dims@ given the flat index @i@.  The resulting list
-- will have the same size as @dims@.
unflattenIndex :: IntegralExp num =>
                  [num] -> num -> [num]
unflattenIndex :: [num] -> num -> [num]
unflattenIndex = [num] -> num -> [num]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndexFromSlices ([num] -> num -> [num])
-> ([num] -> [num]) -> [num] -> num -> [num]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [num] -> [num]
forall a. Int -> [a] -> [a]
drop Int
1 ([num] -> [num]) -> ([num] -> [num]) -> [num] -> [num]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [num] -> [num]
forall num. IntegralExp num => [num] -> [num]
sliceSizes

unflattenIndexFromSlices :: IntegralExp num =>
                            [num] -> num -> [num]
unflattenIndexFromSlices :: [num] -> num -> [num]
unflattenIndexFromSlices [] num
_ = []
unflattenIndexFromSlices (num
size : [num]
slices) num
i =
  (num
i num -> num -> num
forall e. IntegralExp e => e -> e -> e
`quot` num
size) num -> [num] -> [num]
forall a. a -> [a] -> [a]
: [num] -> num -> [num]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndexFromSlices [num]
slices (num
i num -> num -> num
forall a. Num a => a -> a -> a
- (num
i num -> num -> num
forall e. IntegralExp e => e -> e -> e
`quot` num
size) num -> num -> num
forall a. Num a => a -> a -> a
* num
size)

-- | @flattenIndex dims is@ computes the flat index of @is@ into an
-- array with dimensions @dims@.  The length of @dims@ and @is@ must
-- be the same.
flattenIndex :: IntegralExp num =>
                [num] -> [num] -> num
flattenIndex :: [num] -> [num] -> num
flattenIndex [num]
dims [num]
is =
  [num] -> num
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([num] -> num) -> [num] -> num
forall a b. (a -> b) -> a -> b
$ (num -> num -> num) -> [num] -> [num] -> [num]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith num -> num -> num
forall a. Num a => a -> a -> a
(*) [num]
is [num]
slicesizes
  where slicesizes :: [num]
slicesizes = Int -> [num] -> [num]
forall a. Int -> [a] -> [a]
drop Int
1 ([num] -> [num]) -> [num] -> [num]
forall a b. (a -> b) -> a -> b
$ [num] -> [num]
forall num. IntegralExp num => [num] -> [num]
sliceSizes [num]
dims

-- | Given a length @n@ list of dimensions @dims@, @sizeSizes dims@
-- will compute a length @n+1@ list of the size of each possible array
-- slice.  The first element of this list will be the product of
-- @dims@, and the last element will be 1.
sliceSizes :: IntegralExp num =>
              [num] -> [num]
sliceSizes :: [num] -> [num]
sliceSizes [] = [num
1]
sliceSizes (num
n:[num]
ns) =
  [num] -> num
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (num
n num -> [num] -> [num]
forall a. a -> [a] -> [a]
: [num]
ns) num -> [num] -> [num]
forall a. a -> [a] -> [a]
: [num] -> [num]
forall num. IntegralExp num => [num] -> [num]
sliceSizes [num]
ns

{- HLINT ignore sliceSizes -}