dense-linear-algebra: Getting stream fusion to work across `Matrix`'s about dh-core HOT 5 OPEN

Interestingly, this fuses:

testRewrite3 :: Matrix -> Double 
testRewrite3 (Matrix r c v) = U.sum . flip (\(r,c,v) j -> U.generate r (\i -> v `U.unsafeIndex` (j + i * c))) 0 $ (r,c,v)

from dh-core.

It seems that pattern matching is the cause?
This also fuses:

testRewrite4 :: Matrix -> Double 
testRewrite4 m = U.sum . flip (\m1 j -> U.generate (rows m1) (\i -> (_vector m1) `U.unsafeIndex` (j + i * (cols m1)))) 0 $ m

from dh-core.

Bottom line:

column :: Matrix -> Int -> Vector Double
column (Matrix r c v) j= U.generate r (\i -> v `U.unsafeIndex` (j + i * c))
{-# INLINE column #-}

column2 :: Matrix -> Int -> Vector Double
column2 m j = U.generate (rows m) (\i -> (_vector m) `U.unsafeIndex` (j + i * (cols m)))
{-# INLINE column2 #-}

testRewrite0 :: Matrix -> Double
testRewrite0 a = (U.sum . flip column2 0) a

testRewrite1 :: Matrix -> Double
testRewrite1 a = (U.sum . flip column 0) a

testRewrite0 fuses, testRewrite1 doesn't.

However, the fusion is depending on the INLINE and I can't get M.column2 to get inlined across modules (moving the definition of column2 from the main file to Matrix.hs)

from dh-core.

The results are actually a bit disappointing, I'll keep digging

from dh-core.

Back to it! Turns out GHC makes wonders with rewrite rules/inlining. So throwing in for example:

{-# INLINE [1] row #-}
{-# RULES
      "row/fuse"    forall c v r i.  row (M.Matrix r c v) i = U.slice (c*i) c v
  #-}

works amazing.

It gives us some nice stuff:

Say I have matrices A,B, and C = A*B.
Then now I can take the sum of a row of C, without allocating the whole matrix, and a ~10x speedup on my machine.

I would expect to hit a point of diminishing return though, so that's to be explored

from dh-core.