@@ -4,14 +4,17 @@ module Grenade.Layers.Internal.Convolution (
44 , vid2colUnsafe
55 , im2colUnsafe
66 , fittingStarts
7+ , unsafeModifyMatrix
78 ) where
89
9−import Control.Monad.ST ( runST )
10+import Control.Monad.ST ( ST, runST )
1011
11−import Data.STRef ( newSTRef, modifySTRef, writeSTRef, readSTRef )
12+import Data.STRef ( newSTRef, modifySTRef', writeSTRef, readSTRef )
1213import Data.Foldable ( forM_ )
1314import Data.Traversable ( forM )
1415
16+import Foreign.Storable( Storable )
17+
1518import Numeric.LinearAlgebra hiding ( uniformSample, konst )
1619import qualified Numeric.LinearAlgebra.Devel as U
1720
@@ -72,18 +75,17 @@ col2imUnsafe kernelRows kernelColumns strideRows strideColumns destinationRows d
7275
7376 forM_ [0 .. columnMatrixRows - 1] $ \inputRow -> do
7477 inputColumnRef <- newSTRef 0
75− forM_ [0 .. kernelRows -1] $ \kr ->
76− forM_ [0 .. kernelColumns -1] $ \kc -> do
78+ offsetR' <- readSTRef offsetR
79+ offsetC' <- readSTRef offsetC
80+ forM_ [offsetR' .. offsetR' + kernelRows -1] $ \kr ->
81+ forM_ [offsetC' .. offsetC' + kernelColumns -1] $ \kc -> do
7782 inputColumn <- readSTRef inputColumnRef
78− offsetR' <- readSTRef offsetR
79− offsetC' <- readSTRef offsetC
80− U.modifyMatrix dataIm (kr + offsetR') (kc + offsetC') (+ U.atM' columnMatrix inputRow inputColumn)
81− modifySTRef inputColumnRef (+1)
83+ unsafeModifyMatrix dataIm kr kc (+ U.atM' columnMatrix inputRow inputColumn)
84+ modifySTRef' inputColumnRef (+1)
8285
83− offsetC' <- readSTRef offsetC
8486 if offsetC' + kernelColumns < destinationCols
85− then modifySTRef offsetC (+ strideColumns)
86− else writeSTRef offsetC 0 >> modifySTRef offsetR (+ strideRows)
87+ then modifySTRef' offsetC (+ strideColumns)
88+ else writeSTRef offsetC 0 >> modifySTRef' offsetR (+ strideRows)
8789
8890 return dataIm
8991
@@ -104,18 +106,17 @@ col2vidUnsafe kernelRows kernelColumns strideRows strideColumns destinationRows
104106 offsetC <- newSTRef 0
105107 forM_ [0 .. columnMatrixRows - 1] $ \ir -> do
106108 inputColumn <- newSTRef 0
107− forM_ [0 .. kernelRows -1] $ \kr ->
108− forM_ [0 .. kernelColumns -1] $ \kc -> do
109+ offsetR' <- readSTRef offsetR
110+ offsetC' <- readSTRef offsetC
111+ forM_ [offsetR' .. offsetR' + kernelRows -1] $ \kr ->
112+ forM_ [offsetC' .. offsetC' + kernelColumns -1] $ \kc -> do
109113 ic <- readSTRef inputColumn
110− offsetR' <- readSTRef offsetR
111− offsetC' <- readSTRef offsetC
112− U.modifyMatrix dataIm (kr + offsetR') (kc + offsetC') (+ U.atM' columnMatrix ir (ic + offsetM))
113− modifySTRef inputColumn (+1)
114+ unsafeModifyMatrix dataIm kr kc (+ U.atM' columnMatrix ir (ic + offsetM))
115+ modifySTRef' inputColumn (+1)
114116
115− offsetC' <- readSTRef offsetC
116117 if offsetC' + kernelColumns < destinationCols
117− then modifySTRef offsetC (+ strideColumns)
118− else writeSTRef offsetC 0 >> modifySTRef offsetR (+ strideRows)
118+ then modifySTRef' offsetC (+ strideColumns)
119+ else writeSTRef offsetC 0 >> modifySTRef' offsetR (+ strideRows)
119120
120121 U.unsafeFreezeMatrix dataIm
121122
@@ -136,14 +137,14 @@ vid2colUnsafe kernelRows kernelColumns striderows stridecols vidrows vidcols dat
136137 forM_ starts $ \(startRow, startCol) -> do
137138 inputColumnRef <- newSTRef 0
138139 inputRow <- readSTRef inputRowRef
139− forM_ [0 .. kernelRows -1] $ \kr ->
140− forM_ [0 .. kernelColumns -1] $ \kc -> do
140+ forM_ [startRow .. startRow + kernelRows -1] $ \kr ->
141+ forM_ [startCol .. startCol + kernelColumns -1] $ \kc -> do
141142 inputColumn <- readSTRef inputColumnRef
142− U.modifyMatrix dataCol inputRow (inputColumn + offsetC') (+ U.atM' dataIm (kr + startRow) (kc + startCol))
143− modifySTRef inputColumnRef (+1)
144− modifySTRef inputRowRef (+1)
143+ unsafeModifyMatrix dataCol inputRow (inputColumn + offsetC') (+ U.atM' dataIm kr kc)
144+ modifySTRef' inputColumnRef (+1)
145+ modifySTRef' inputRowRef (+1)
145146
146− modifySTRef offsetC (+ kernelSize)
147+ modifySTRef' offsetC (+ kernelSize)
147148
148149 return dataCol
149150
@@ -159,15 +160,19 @@ im2colUnsafe kernelRows kernelColumns striderows stridecols dataIm = U.runSTMatr
159160 forM_ starts $ \(startRow, startCol) -> do
160161 inputColumnRef <- newSTRef 0
161162 inputRow <- readSTRef inputRowRef
162− forM_ [0 .. kernelRows -1] $ \kr ->
163− forM_ [0 .. kernelColumns -1] $ \kc -> do
163+ forM_ [startRow .. startRow + kernelRows -1] $ \kr ->
164+ forM_ [startCol .. startCol + kernelColumns -1] $ \kc -> do
164165 inputColumn <- readSTRef inputColumnRef
165− U.modifyMatrix dataCol inputRow inputColumn (+ U.atM' dataIm (kr + startRow) (kc + startCol))
166− modifySTRef inputColumnRef (+1)
167− modifySTRef inputRowRef (+1)
166+ unsafeModifyMatrix dataCol inputRow inputColumn (+ U.atM' dataIm kr kc)
167+ modifySTRef' inputColumnRef (+1)
168+ modifySTRef' inputRowRef (+1)
168169
169170 return dataCol
170171
172+unsafeModifyMatrix :: (Storable t) => U.STMatrix s t -> Int -> Int -> (t -> t) -> ST s ()
173+unsafeModifyMatrix x r c f = U.unsafeReadMatrix x r c >>= U.unsafeWriteMatrix x r c . f
174+{-# INLINE unsafeModifyMatrix #-}
175+
171176
172177-- | Returns the starting sub matrix locations which fit inside the larger matrix for the
173178-- convolution. Takes into account the stride and kernel size.