Stencil with self dependency and waterfall parallelization

[lehins]

Recently there was a request on dataHaskell gitter about creating an array in parallel, where generation of an array element depends on some of its neighbors. Here is a solution for 2D matrices:

waterfallCreate ::
     Mutable r Ix2 a => Sz2 -> (Maybe a -> Maybe a -> a) -> (a -> a -> a) -> Array r Ix2 a
waterfallCreate sz@(Sz2 m n) g f =
  unsafePerformIO $
  createArray_ Par sz $ \scheduler marr -> do
    void $ write marr (0 :. 0) $ g Nothing Nothing
    forM_ (1 ..: m) $ \i ->
      unsafeWrite marr (i :. 0) . g Nothing . Just =<< unsafeRead marr (i - 1 :. 0)
    forM_ (1 ..: n) $ \j -> do
      unsafeWrite marr (0 :. j) . (`g` Nothing) . Just  =<< unsafeRead marr (0 :. j - 1)
      let jIter = rangeStepSize Seq j (-1) (Sz (min (m - 1) j))
      iforSchedulerM_ scheduler jIter $ \i' -> writeF marr (i' + 1)
    forM_ (2 ..: m) $ \i -> do
      let jIter = rangeStepSize Seq (n - 1) (-1) (Sz (min (n - 1) (m - i)))
      iforSchedulerM_ scheduler jIter $ \i' -> writeF marr (i' + i)
  where
    writeF marr i j = do
      left <- unsafeRead marr (i :. j - 1)
      top <- unsafeRead marr (i - 1 :. j)
      unsafeWrite marr (i :. j) $ f left top
    {-# INLINE writeF #-}
{-# INLINE waterfallCreate #-}

Above solution is a good demonstration of the problem, but the full solution will require:

• support for arbitrary dimensions
• ability to select which neighbors there is a dependency on
• implemented as a stencil, not only for creation of new arrays
• use same optimizations as in DW array
• block wise parallelization along the diagonal, instead of single element diagonal.