1 files changed, 118 insertions, 13 deletions

diff --git a/packages/base/src/Internal/ST.hs b/packages/base/src/Internal/ST.hs
index 7d54e6d..326b90a 100644
--- a/packages/base/src/Internal/ST.hs
+++ b/packages/base/src/Internal/ST.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE Rank2Types    #-}
 {-# LANGUAGE BangPatterns  #-}
 {-# LANGUAGE ViewPatterns #-}
+{-# LANGUAGE PatternSynonyms #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -30,14 +31,20 @@ module Internal.ST (
     unsafeThawVector, unsafeFreezeVector,
     newUndefinedMatrix,
     unsafeReadMatrix, unsafeWriteMatrix,
-    unsafeThawMatrix, unsafeFreezeMatrix
+    unsafeThawMatrix, unsafeFreezeMatrix,
+    setRect
 ) where
 
 import Internal.Vector
 import Internal.Matrix
 import Internal.Vectorized
+import Internal.Devel ((#|))
 import Control.Monad.ST(ST, runST)
-import Foreign.Storable(Storable, peekElemOff, pokeElemOff)
+import Control.Monad
+import Data.Function
+import Data.Int
+import Foreign.Ptr
+import Foreign.Storable
 import Control.Monad.ST.Unsafe(unsafeIOToST)
 
 {-# INLINE ioReadV #-}
@@ -121,7 +128,7 @@ ioWriteM m r c val = ioWriteV (xdat m)  (r * xRow m + c * xCol m) val
 
 newtype STMatrix s t = STMatrix (Matrix t)
 
-thawMatrix :: Element t => Matrix t -> ST s (STMatrix s t)
+thawMatrix :: Storable t => Matrix t -> ST s (STMatrix s t)
 thawMatrix = unsafeIOToST . fmap STMatrix . cloneMatrix
 
 unsafeThawMatrix :: Storable t => Matrix t -> ST s (STMatrix s t)
@@ -142,17 +149,17 @@ unsafeWriteMatrix  (STMatrix x) r c = unsafeIOToST . ioWriteM x r c
 modifyMatrix :: (Storable t) => STMatrix s t -> Int -> Int -> (t -> t) -> ST s ()
 modifyMatrix x r c f = readMatrix x r c >>= return . f >>= unsafeWriteMatrix x r c
 
-liftSTMatrix :: (Element t) => (Matrix t -> a) -> STMatrix s t -> ST s a
+liftSTMatrix :: (Storable t) => (Matrix t -> a) -> STMatrix s t -> ST s a
 liftSTMatrix f (STMatrix x) = unsafeIOToST . fmap f . cloneMatrix $ x
 
 unsafeFreezeMatrix :: (Storable t) => STMatrix s t -> ST s (Matrix t)
 unsafeFreezeMatrix (STMatrix x) = unsafeIOToST . return $ x
 
 
-freezeMatrix :: (Element t) => STMatrix s t -> ST s (Matrix t)
+freezeMatrix :: (Storable t) => STMatrix s t -> ST s (Matrix t)
 freezeMatrix m = liftSTMatrix id m
 
-cloneMatrix :: Element t => Matrix t -> IO (Matrix t)
+cloneMatrix :: Storable t => Matrix t -> IO (Matrix t)
 cloneMatrix m = copy (orderOf m) m
 
 {-# INLINE safeIndexM #-}
@@ -172,7 +179,7 @@ readMatrix = safeIndexM unsafeReadMatrix
 writeMatrix :: Storable t => STMatrix s t -> Int -> Int -> t -> ST s ()
 writeMatrix = safeIndexM unsafeWriteMatrix
 
-setMatrix :: Element t => STMatrix s t -> Int -> Int -> Matrix t -> ST s ()
+setMatrix :: Storable t => STMatrix s t -> Int -> Int -> Matrix t -> ST s ()
 setMatrix (STMatrix x) i j m = unsafeIOToST $ setRect i j m x
 
 newUndefinedMatrix :: Storable t => MatrixOrder -> Int -> Int -> ST s (STMatrix s t)
@@ -210,7 +217,7 @@ data RowOper t = AXPY t Int Int  ColRange
                | SCAL t RowRange ColRange
                | SWAP Int Int    ColRange
 
-rowOper :: (Num t, Element t) => RowOper t -> STMatrix s t -> ST s ()
+rowOper :: (Num t, Storable t) => RowOper t -> STMatrix s t -> ST s ()
 
 rowOper (AXPY x i1 i2 r) (STMatrix m) = unsafeIOToST $ rowOp 0 x i1' i2' j1 j2 m
   where
@@ -230,8 +237,8 @@ rowOper (SWAP i1 i2 r) (STMatrix m) = unsafeIOToST $ rowOp 2 0 i1' i2' j1 j2 m
     i2' = i2 `mod` (rows m)
 
 
-extractMatrix :: Element a => STMatrix t a -> RowRange -> ColRange -> ST s (Matrix a)
-extractMatrix (STMatrix m) rr rc = unsafeIOToST (extractR (orderOf m) m 0 (idxs[i1,i2]) 0 (idxs[j1,j2]))
+extractMatrix :: Storable a => STMatrix t a -> RowRange -> ColRange -> ST s (Matrix a)
+extractMatrix (STMatrix m) rr rc = unsafeIOToST (extractAux (orderOf m) m 0 (idxs[i1,i2]) 0 (idxs[j1,j2]))
   where
     (i1,i2) = getRowRange (rows m) rr
     (j1,j2) = getColRange (cols m) rc
@@ -239,19 +246,117 @@ extractMatrix (STMatrix m) rr rc = unsafeIOToST (extractR (orderOf m) m 0 (idxs[
 -- | r0 c0 height width
 data Slice s t = Slice (STMatrix s t) Int Int Int Int
 
-slice :: Element a => Slice t a -> Matrix a
+slice :: Storable a => Slice t a -> Matrix a
 slice (Slice (STMatrix m) r0 c0 nr nc) = subMatrix (r0,c0) (nr,nc) m
 
-gemmm :: Element t => t -> Slice s t -> t -> Slice s t -> Slice s t -> ST s ()
+gemmm :: (Storable t, Num t) => t -> Slice s t -> t -> Slice s t -> Slice s t -> ST s ()
 gemmm beta (slice->r) alpha (slice->a) (slice->b) = res
   where
     res = unsafeIOToST (gemm v a b r)
     v = fromList [alpha,beta]
 
 
-mutable :: Element t => (forall s . (Int, Int) -> STMatrix s t -> ST s u) -> Matrix t -> (Matrix t,u)
+mutable :: Storable t => (forall s . (Int, Int) -> STMatrix s t -> ST s u) -> Matrix t -> (Matrix t,u)
 mutable f a = runST $ do
    x <- thawMatrix a
    info <- f (rows a, cols a) x
    r <- unsafeFreezeMatrix x
    return (r,info)
+
+
+
+setRect :: Storable t => Int -> Int -> Matrix t -> Matrix t -> IO ()
+setRect i j m r = (m Internal.Matrix.#! r) (setRectStorable (fi i) (fi j)) #|"setRect"
+
+setRectStorable :: Storable t =>
+                    Int32 -> Int32
+                    -> Int32 -> Int32 -> Int32 -> Int32 -> {- const -} Ptr t
+                    -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr t
+                    -> IO Int32
+setRectStorable i j mr mc mXr mXc mp rr rc rXr rXc rp = do
+    ($ 0) $ fix $ \aloop a -> when (a<mr) $ do
+        ($ 0) $ fix $ \bloop b -> when (b<mc) $ do
+            let x = a+i
+                y = b+j
+            when (0<=x && x<rr && 0<=y && y<rc) $ do
+                pokeElemOff rp (fromIntegral $ rXr*x + rXc*y)
+                    =<< peekElemOff mp (fromIntegral $ mXr*a + mXc*b)
+            bloop (succ b)
+        aloop (succ a)
+    return 0
+
+rowOp :: (Storable t, Num t) => Int -> t -> Int -> Int -> Int -> Int -> Matrix t -> IO ()
+rowOp = rowOpAux rowOpStorable
+
+pattern BAD_CODE = 2001
+
+rowOpStorable :: (Storable t, Num t) =>
+                 Int32 -> Ptr t -> Int32 -> Int32 -> Int32 -> Int32
+                 -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr t
+                 -> IO Int32
+rowOpStorable 0 pa i1 i2 j1 j2 rr rc rXr rXc rp = do
+    -- AXPY_IMP
+    a <- peek pa
+    ($ j1) $ fix $ \jloop j -> when (j<=j2) $ do
+        ri1j <- peekElemOff rp $ fromIntegral $ rXr*i1 + rXc*j
+        let i2j = fromIntegral $ rXr*i2 + rXc*j
+        ri2j <- peekElemOff rp i2j
+        pokeElemOff rp i2j $ ri2j + a*ri1j
+        jloop (succ j)
+    return 0
+rowOpStorable 1 pa i1 i2 j1 j2 rr rc rXr rXc rp = do
+    -- SCAL_IMP
+    a <- peek pa
+    ($ i1) $ fix $ \iloop i -> when (i<=i2) $ do
+        ($ j1) $ fix $ \jloop j -> when (j<=j2) $ do
+            let rijp = rp `plusPtr` fromIntegral (rXr*i + rXc*j)
+            rij <- peek rijp
+            poke rijp $ a * rij
+            jloop (succ j)
+        iloop (succ i)
+    return 0
+rowOpStorable 2 pa i1 i2 j1 j2 rr rc rXr rXc rp | i1 == i2 = return 0
+rowOpStorable 2 pa i1 i2 j1 j2 rr rc rXr rXc rp = do
+    -- SWAP_IMP
+    ($ j1) $ fix $ \kloop k -> when (k<=j2) $ do
+        let i1k = fromIntegral $ rXr*i1 + rXc*k
+            i2k = fromIntegral $ rXr*i2 + rXc*k
+        aux <- peekElemOff rp i1k
+        pokeElemOff rp i1k =<< peekElemOff rp i2k
+        pokeElemOff rp i2k aux
+        kloop (succ k)
+    return 0
+rowOpStorable _ pa i1 i2 j1 j2 rr rc rXr rXc rp = do
+    return BAD_CODE
+
+gemm :: (Storable t, Num t) => Vector t -> Matrix t -> Matrix t -> Matrix t -> IO ()
+gemm v m1 m2 m3 = (v Internal.Matrix.# m1 Internal.Matrix.# m2 Internal.Matrix.#! m3) gemmStorable #|"gemm"
+
+-- ScalarLike t
+gemmStorable :: (Storable t, Num t) =>
+    Int32 -> Ptr t -- VECG(T,c)
+    -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr t -- MATG(T,a)
+    -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr t -- MATG(T,b)
+    -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr t -- MATG(T,r)
+    -> IO Int32
+gemmStorable cn cp
+    ar ac aXr aXc ap
+    br bc bXr bXc bp
+    rr rc rXr rXc rp = do
+        a <- peek cp
+        b <- peekElemOff cp 1
+        ($ 0) $ fix $ \iloop i -> when (i<rr) $ do
+            ($ 0) $ fix $ \jloop j -> when (j<rc) $ do
+                let kloop k !t fin
+                        | k<ac = do
+                            aik <- peekElemOff ap (fromIntegral $ i*aXr + k*aXc)
+                            bkj <- peekElemOff bp (fromIntegral $ k*bXr + j*bXc)
+                            kloop (succ k) (t + aik*bkj) fin
+                        | otherwise = fin t
+                kloop 0 0 $ \t -> do
+                    let ij = fromIntegral $ i*rXr + j*rXc
+                    rij <- peekElemOff rp ij
+                    pokeElemOff rp ij (b*rij + a*t)
+                jloop (succ j)
+            iloop (succ i)
+        return 0