{-# LANGUAGE FlexibleContexts #-}
module Stencil2 where
import Control.Monad
import Control.Exception
import System.Random.MWC
import Data.Array.Unboxed hiding (Array)
import Data.Array.Accelerate hiding (round, min, max, fromIntegral)
import qualified Data.Array.IArray as IArray
stencil2D2 :: Floating (Exp a) => Stencil3x3 a -> Stencil3x3 a -> Exp a
stencil2D2 ((_,t,_), (_,x,_), (_,b,_))
((_,_,_), (l,y,r), (_,_,_)) = t + b + l + r - ((x+y) / 2)
test_stencil2_2D :: Int -> IO (() -> UArray (Int,Int) Float, () -> Acc (Array DIM2 Float))
test_stencil2_2D n2 = withSystemRandom $ \gen -> do
let n = round $ sqrt (fromIntegral n2 :: Double)
m = n * 2
u = m `div` 3
v = n + m
m1 <- listArray ((0,0),(n-1,m-1)) `fmap` replicateM (n*m) (uniformR (-1,1) gen) :: IO (UArray (Int,Int) Float)
m2 <- listArray ((0,0),(u-1,v-1)) `fmap` replicateM (u*v) (uniformR (-1,1) gen) :: IO (UArray (Int,Int) Float)
m1' <- let m1' = fromIArray m1 in evaluate (m1' `indexArray` (Z:.0:.0)) >> return m1'
m2' <- let m2' = fromIArray m2 in evaluate (m2' `indexArray` (Z:.0:.0)) >> return m2'
--
return (\() -> run_ref m1 m2, \() -> run_acc m1' m2')
where
run_acc xs ys = stencil2 stencil2D2 Mirror (use xs) Wrap (use ys)
run_ref xs ys =
let (_,(n,m)) = bounds xs
(_,(u,v)) = bounds ys
sh = ((0,0), (n `min` u, m `min` v))
-- boundary conditions are placed on the *source* arrays
--
get1 (x,y) = xs IArray.! (mirror n x, mirror m y)
get2 (x,y) = ys IArray.! (wrap u x, wrap v y)
mirror sz i
| i < 0 = -i
| i > sz = sz - (i-sz)
| otherwise = i
wrap sz i
| i < 0 = sz + i + 1
| i > sz = i - sz - 1
| otherwise = i
f (ix,iy) = let t = get1 (ix, iy-1)
b = get1 (ix, iy+1)
x = get1 (ix, iy)
l = get2 (ix-1,iy)
r = get2 (ix+1,iy)
y = get2 (ix, iy)
in
t + b + l + r - ((x+y) / 2)
in
array sh [(ix, f ix) | ix <- range sh]
-- Main
-- ----
run2D :: String -> Int -> IO (() -> UArray (Int,Int) Float, () -> Acc (Array DIM2 Float))
run2D "2D" = test_stencil2_2D
run2D x = error $ "unknown variant: " ++ x