More flexible indexing

[tmcdonell]

Here's something that I have been playing with recently. The details & consequences are not fully fleshed out, but I include a small sample implementation to give the flavour. The core is to change the type of array indexes to:

type DIM0 = Z
type DIM1 = Int
type DIM2 = DIM1 :. Int     -- Int :. Int
type DIM3 = DIM2 :. Int     -- Int :. Int :. Int
type DIM4 = DIM3 :. Int     -- Int :. Int :. Int :. Int

so that for DIM2 and higher, they are open at both the left and the right. We then have (pattern synonyms) (:>) and (:<) to add/remove an index on the right (as we have now) as well as the left, respectively. There are a few examples at the bottom of the attached.

Comments/suggestions/criticisms welcome!

Example implementation

{-# LANGUAGE AllowAmbiguousTypes  #-}
{-# LANGUAGE ConstraintKinds      #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE PatternSynonyms      #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns         #-}

import GHC.TypeNats
import Text.Printf
import Data.Vector                              ( Vector )
import qualified Data.Vector                    as V


data Z = Z
  deriving Show

infixl 3 :.
data tail :. head = tail :. head

type DIM0 = Z
type DIM1 = Int
type DIM2 = DIM1 :. Int     -- Int :. Int
type DIM3 = DIM2 :. Int     -- Int :. Int :. Int
type DIM4 = DIM3 :. Int     -- Int :. Int :. Int :. Int

instance (Show sh, Show sz) => Show (sh :. sz) where
  showsPrec p (sh :. sz) =
    showsPrec p sh . showString " :. " . showsPrec p sz

type family Lower sh where
  Lower Int         = Z
  Lower (sh :. Int) = sh

type family Higher sh where
  Higher Z  = Int
  Higher sh = sh  :. Int

data ShapeR sh where
  ShapeRz    :: ShapeR Z
  ShapeRint  :: ShapeR Int
  ShapeRsnoc :: 1 <= Rank sh => ShapeR sh -> ShapeR (sh :. Int)

-- This is equivalent to (:.), but drops the Z constructor for
-- non-singleton dimensions
--
infixl 3 :>
pattern (:>) :: (Shape sh, Shape (Lower sh), Higher (Lower sh) ~ sh, 1 <= Rank sh) => Lower sh -> Int -> sh
pattern (:>) sh sz <- (unsnoc -> (sh, sz))
  where
    (:>) = snoc

infixr 3 :<
pattern (:<) :: (Shape sh, Shape (Lower sh), Higher (Lower sh) ~ sh, 1 <= Rank sh, 1 <= Rank (Lower sh)) => Int -> Lower sh -> sh
pattern (:<) sz sh <- (uncons -> (sz, sh))
  where
    (:<) = cons

class Shape sh where
  type Rank sh :: Nat
  shapeR :: ShapeR sh

instance Shape Z where
  type Rank Z = 0
  shapeR = ShapeRz

instance Shape Int where
  type Rank Int = 1
  shapeR = ShapeRint

instance (Shape sh, 1 <= Rank sh) => Shape (sh :. Int) where
  type Rank (sh :. Int) = 1 + Rank sh
  shapeR = ShapeRsnoc shapeR


-- Add a new _innermost_-dimension on the _right_
--
snoc :: forall sh. Shape sh => sh -> Int -> Higher sh
snoc sh i =
  case shapeR @sh of
    ShapeRz      -> i
    ShapeRint    -> sh :. i
    ShapeRsnoc _ -> sh :. i

-- Remove an _innermost_-dimension from the _right_
--
unsnoc :: forall sh. (Shape sh, 1 <= Rank sh) => sh -> (Lower sh, Int)
unsnoc sh =
  case shapeR @sh of
    ShapeRint    -> (Z, sh)
    ShapeRsnoc _ -> let t :. h = sh in (t, h)

-- Add a new _outermost_-dimension on the _left_
--
cons :: (Shape sh, 1 <= Rank sh) => Int -> sh -> Higher sh
cons = cons' shapeR
  where
    cons' :: 1 <= Rank sh => ShapeR sh -> Int -> sh -> Higher sh
    cons' shR i sh =
      case shR of
        ShapeRint                    -> i :. sh
        ShapeRsnoc ShapeRint         -> let t :. h = sh in i :. t :. h
        ShapeRsnoc shR'@ShapeRsnoc{} -> let t :. h = sh
                                            sh'    = cons' shR' i t
                                         in
                                         sh' :. h

-- Remove an _outermost_-dimension from the _left_
--
uncons :: (Shape sh, 1 <= Rank sh) => sh -> (Int, Lower sh)
uncons = uncons' shapeR
  where
    uncons' :: 1 <= Rank sh => ShapeR sh -> sh -> (Int, Lower sh)
    uncons' shR sh =
      case shR of
        ShapeRint                    -> (sh, Z)
        ShapeRsnoc ShapeRint         -> let t :. h = sh in (t,h)
        ShapeRsnoc shR'@ShapeRsnoc{} -> let t :. h  = sh
                                            (i, t') = uncons' shR' t
                                         in
                                         (i, t' :. h)

toIndex :: Shape sh => sh -> sh -> Int
toIndex = go shapeR
  where
    go :: ShapeR sh -> sh -> sh -> Int
    go ShapeRz          Z          Z         = 0
    go ShapeRint        _          i         = i
    go (ShapeRsnoc shR) (sh :. sz) (ix :. i) = go shR sh ix * sz + i

fromIndex :: Shape sh => sh -> Int -> sh
fromIndex = go shapeR
  where
    go :: ShapeR sh -> sh -> Int -> sh
    go ShapeRz          Z          _ = Z
    go ShapeRint        _          i = i
    go (ShapeRsnoc shR) (sh :. sz) i = go shR sh q :. r
      where
        (q,r) = quotRem i sz

size :: Shape sh => sh -> Int
size = go shapeR
  where
    go :: ShapeR sh -> sh -> Int
    go ShapeRz          Z          = 1
    go ShapeRint        x          = x
    go (ShapeRsnoc shR) (sh :. sz) = go shR sh * sz

data Array sh e where
  Array :: { shape   :: sh
           , payload :: Vector e
           } -> Array sh e

instance (Shape sh, Show sh, Show e) => Show (Array sh e) where
  show (Array sh adata) =
    let sh' :: String
        sh' = case shapeR @sh of
                ShapeRsnoc{} -> "(" ++ show sh ++ ")"
                _            -> show sh
        adata' = show (V.toList adata)
    in
    printf "Array %s %s" sh' adata'


data All = All

data ReduceR ix slice full where
  ReduceRz    :: ReduceR All Z   Int  -- outermost-dimension is reduced
  ReduceRint  :: ReduceR Int Int Int  -- outermost-dimension is kept
  ReduceRall  :: 1 <= Rank full => ShapeR slice -> ShapeR full -> ReduceR ix slice full -> ReduceR (ix :. All) slice          (Higher full) -- inner dimension is reduced
  ReduceRnext :: 1 <= Rank full => ShapeR slice -> ShapeR full -> ReduceR ix slice full -> ReduceR (ix :. Int) (Higher slice) (Higher full) -- inner dimension is kept

instance Show (ReduceR ix slice full) where
  show ReduceRz               = "ReduceRz"
  show ReduceRint             = "ReduceRint"
  show (ReduceRall _ _ rest)  = "ReduceRall (" ++ show rest ++ ")"
  show (ReduceRnext _ _ rest) = "ReduceRnext (" ++ show rest ++ ")"

class (Shape (SliceShape sh), Shape (FullShape sh)) => Reduce sh where
  type SliceShape sh
  type FullShape  sh
  reduceIndex :: ReduceR sh (SliceShape sh) (FullShape sh)

instance Reduce All where
  type SliceShape All = Z
  type FullShape  All = Int
  reduceIndex = ReduceRz

instance Reduce Int where
  type SliceShape Int = Int
  type FullShape  Int = Int
  reduceIndex = ReduceRint

instance (Reduce sl, Shape (Higher (FullShape sl)), 1 <= Rank (FullShape sl)) => Reduce (sl :. All) where
  type SliceShape (sl :. All) = SliceShape sl
  type FullShape  (sl :. All) = Higher (FullShape sl)
  reduceIndex = ReduceRall shapeR shapeR reduceIndex

instance (Reduce sl, Shape (Higher (SliceShape sl)), Shape (Higher (FullShape sl)), 1 <= Rank (FullShape sl)) => Reduce (sl :. Int) where
  type SliceShape (sl :. Int) = Higher (SliceShape sl)
  type FullShape  (sl :. Int) = Higher (FullShape sl)
  reduceIndex = ReduceRnext shapeR shapeR reduceIndex

restrict
    :: ReduceR slix slice full
    -> full
    -> slice
restrict ReduceRz   _  = Z
restrict ReduceRint sz = sz
restrict (ReduceRall _ fullR reduceR) full =
  case fullR of
    ShapeRint    -> let sl :. _ = full in restrict reduceR sl
    ShapeRsnoc{} -> let sl :. _ = full in restrict reduceR sl
restrict (ReduceRnext sliceR fullR reduceR) full =
  case fullR of
    ShapeRint   -> let sl :. sh = full
                    in case sliceR of
                         ShapeRz      -> sh
                         ShapeRint    -> restrict reduceR sl :. sh
                         ShapeRsnoc{} -> restrict reduceR sl :. sh
    ShapeRsnoc{} -> let sl :. sh = full
                    in case sliceR of
                         ShapeRz      -> sh
                         ShapeRint    -> restrict reduceR sl :. sh
                         ShapeRsnoc{} -> restrict reduceR sl :. sh

iter
    :: ReduceR slix slice full
    -> slice
    -> full
    -> (full -> e)
    -> (e -> e -> e)
    -> e
    -> e
iter ReduceRz Z sz f c r =
  let go i | i >= sz   = r
           | otherwise = f i `c` go (i+1)
   in go 0
iter ReduceRint                   ix    _    f _ _ = f ix
iter (ReduceRall _ fullR reduceR) slice full f c r =
  case fullR of
    ShapeRint    -> let sh :. sz = full
                        go (ix :. i) | i >= sz   = r
                                     | otherwise = f (ix :. i) `c` go (ix :. i+1)
                     in
                     iter reduceR slice sh (\ix -> go (ix :. 0)) c r

    ShapeRsnoc{} -> let sh :. sz = full
                        go (ix :. i) | i >= sz   = r
                                     | otherwise = f (ix :. i) `c` go (ix :. i+1)
                     in
                     iter reduceR slice sh (\ix -> go (ix :. 0)) c r
iter (ReduceRnext sliceR fullR reduceR) slice full f c r =
  case fullR of
    ShapeRint   -> let sh :. _ = full
                    in case sliceR of
                         ShapeRz      -> iter reduceR Z sh (\ix -> f (ix :. slice)) c r
                         ShapeRint    -> let sl :. i = slice in iter reduceR sl sh (\ix -> f (ix :. i)) c r
                         ShapeRsnoc{} -> let sl :. i = slice in iter reduceR sl sh (\ix -> f (ix :. i)) c r

    ShapeRsnoc{} -> let sh :. _ = full
                    in case sliceR of
                         ShapeRz      -> iter reduceR Z sh (\ix -> f (ix :. slice)) c r
                         ShapeRint    -> let sl :. i = slice in iter reduceR sl sh (\ix -> f (ix :. i)) c r
                         ShapeRsnoc{} -> let sl :. i = slice in iter reduceR sl sh (\ix -> f (ix :. i)) c r


-- Prelude
-- -------

generate :: Shape sh => sh -> (sh -> e) -> Array sh e
generate sh f =
  Array sh $ V.generate (size sh) (f . fromIndex sh)

infixl 9 !
(!) :: Shape sh => Array sh e -> sh -> e
Array sh adata ! ix = adata V.! toIndex sh ix


-- Reduce an array along arbitrary dimensions. Here 'All' specifies
-- a dimension which should be reduced all. For example, given `f` and `z`
-- and input `xs :: Array DIM4 Int`, then:
--
--  * reduce @(Int :. Int :. Int :. All) f z :: Array DIM3 Int
--      ...is equivalent to our existing reduction along the inner-most dimension
--
--  * reduce @(All :. Int :. Int :. Int) f z :: Array DIM3 Int
--      ...reduces along the outer-most dimension
--
--  * reduce @(All :. Int :. All :. Int) f z :: Array DIM2 Int
--      ...reduces two _non-continuous_ dimensions, however useful that might be
--
reduce :: forall slix e. Reduce slix
       => (e -> e -> e)
       -> e
       -> Array (FullShape  slix) e
       -> Array (SliceShape slix) e
reduce f z arr =
  let reduceR = reduceIndex @slix
      sh      = shape arr
      sh'     = restrict reduceR sh
   in generate sh' (\ix -> iter reduceR ix sh (arr!) f z)


-- Examples
-- --------

arr3 :: Array DIM3 Int
arr3 = let sh = 2 :> 3 :> 4
        in generate sh (toIndex sh)

arr2 :: Array DIM2 Int
arr2 = let sh = 2 :> 5
        in generate sh (toIndex sh)

arr1 :: Array DIM1 Int
arr1 = generate 10 (toIndex 10)

t1 :: Array DIM2 Int
t1 = let _ :< sh = shape arr3
      in generate sh (\ix -> arr3 ! (1 :< ix))

t2 :: Array DIM1 Int
t2 = let _ :> sh :> _ = shape arr3
      in generate sh (\ix -> arr3 ! (0 :> ix :> 2))

t3 :: Array DIM2 Int
t3 = reduce @(Int :. Int :. All) (+) 0 arr3

t4 :: Array DIM2 Int
t4 = reduce @(All :. Int :. Int) (+) 0 arr3

t5 :: Array DIM1 Int
t5 = reduce @(Int :. All :. All) (+) 0 arr3

t6 :: Array DIM1 Int
t6 = reduce @(All :. Int :. All) (+) 0 arr3

t7 :: Array DIM0 Int
t7 = reduce @All (+) 0 arr1

t8 :: Array DIM1 Int
t8 = reduce @(Int :. All) (+) 0 arr2

t9 :: Array DIM1 Int
t9 = reduce @(All :. Int) (+) 0 arr2

[dpvanbalen]

I'm not sure I quite understand why we need the first part (type DIM1 = Int) to get the second part (pattern synonyms (:>) and (:<))? In the example implementation, you recurse the shape like this:

let t :. h = sh
    sh'    = cons' shR' i t
in sh' :. h

In the current way, with type DIM1 = Z :. Int, couldn't you achieve the same thing? I believe I even saw implementations of such a cons function in Accelerate internals before, is there a reason why we can't currently export a neat interface to those?

[tmcdonell]

Yes you are right, we don't really need to drop the left-hand Z. It just felt nicer to think about it as being open on both ends, rather than inserting something in the last position of the list. A more complicated structure than (:.) would allow us to do it directly (ala Data.Sequence), but I don't think that's worth it.

[tmcdonell]

If you ignore how it's actually implemented, my feeling was that:

5 :< (3 :. 2)  ==>  5 :. 3 :. 2

feels better than

5 :< (Z :. 3 :. 2)  ==>  5 :. Z :. 3 :. 2  ==>  Z :. 5 :. 3 :. 2

But of course this is all open for discussion...

[ivogabe]

I agree with David that it would make sense to keep the Z on the left, it seems that we would have less edge-cases by always including the Z. It would be nice to keep this mainly a frontend-feature, without making the internals more complex. We should still be able to shift the Z when adding a new dimension on the left with :< with type families.

[dpvanbalen]

The proposed way definitely looks cleaner, I suppose it's more of a tradeoff between this cleanness and some internal complexity. To avoid changing the internals, you could, I think, even make a constructor like (:..) (is that already taken for lifted shapes?), which would mean:
type n :.. m = Z :. n :. m
This way you'd get:
4 :< (3 :.. 2 :. 1) == 4 :< (Z :. 3 :. 2 :. 1) == Z :. 4 :. 3 :. 2 :. 1 == 4 :.. 3 :. 2 :. 1
Correct me if this is not possible or only more confusing! You could argue about which version the Show instance should represent.

That said, looking at this again, the 'complexity' of the internals really shouldn't be that high with this new approach. The more I look at it, the more I'm in favour.

[tmcdonell]

Yes I don't think the representation types would need to change, this is just a front-end feature.