Ticket #3736: Main.hs

File Main.hs, 5.9 KB (added by igloo, 6 years ago)
Line 
1{-# OPTIONS_GHC -funbox-strict-fields -O #-}
2{-# LANGUAGE ExistentialQuantification #-}
3
4{- OPTIONS_GHC -ddump-simpl -ddump-asm -}
5
6module Main (main) where
7
8import GHC.Float (float2Int, int2Float)
9
10import System.Environment
11
12import Prelude hiding           (null
13                                ,lines,unlines
14                                ,writeFile
15                                )
16
17import Control.Exception        (assert, bracket, )
18
19import Foreign.Marshal.Array    (advancePtr)
20import Foreign.Ptr              (minusPtr)
21import Foreign.Storable         (Storable(..))
22
23import Control.Monad            (when)
24
25import System.IO                (openBinaryFile, hClose,
26                                 hPutBuf,
27                                 Handle, IOMode(..))
28
29import System.IO.Unsafe         (unsafePerformIO)
30
31import Foreign.Ptr              (Ptr)
32import Foreign.ForeignPtr       (ForeignPtr, withForeignPtr, )
33import Foreign.Marshal.Array    (copyArray)
34
35import qualified Foreign.ForeignPtr as F
36
37type Phase = (Float, Float, Float)
38
39{-# INLINE saw #-}
40saw :: Num a => a -> a
41saw t = 1-2*t
42
43{-# INLINE fraction #-}
44fraction :: Float -> Float
45fraction x = x - int2Float (float2Int x)
46
47{-# INLINE generator0Freq #-}
48generator0Freq :: Float -> Float -> Maybe (Float, Float)
49generator0Freq freq =
50   \p -> Just (saw p, fraction (p+freq))
51
52infixl 6 `mix`, `mixGen`
53
54{-# INLINE mix #-}
55mix ::
56   (Num y) =>
57   (s -> Maybe (y, s)) ->
58   (t -> Maybe (y, t)) ->
59   ((s,t) -> Maybe (y, (s,t)))
60mix f g (s0,t0) =
61   do (a,s1) <- f s0
62      (b,t1) <- g t0
63      return ((a+b), (s1,t1))
64
65data Generator a =
66   forall s.
67      Generator (s -> Maybe (a, s)) s
68
69{-# INLINE runGeneratorMonolithic #-}
70runGeneratorMonolithic :: Int -> Generator Float -> Vector Float
71runGeneratorMonolithic size' (Generator f s) =
72   fst $ unfoldrN size' f s
73
74{- SPECIALISE INLINE generator0Gen :: Float -> Float -> Generator Float -}
75{-# INLINE generator0Gen #-}
76generator0Gen :: Float -> Float -> Generator Float
77generator0Gen freq phase =
78   Generator (\p -> Just (saw p, fraction (p+freq))) phase
79
80{- SPECIALISE INLINE mixGen :: Generator Float -> Generator Float -> Generator Float -}
81{-# INLINE mixGen #-}
82mixGen ::
83   (Num y) =>
84   Generator y ->
85   Generator y ->
86   Generator y
87mixGen (Generator f s) (Generator g t) =
88   Generator (\(s0,t0) ->
89      do (a,s1) <- f s0
90         (b,t1) <- g t0
91         return ((a+b), (s1,t1))) (s,t)
92
93{-# INLINE dl #-}
94dl :: Phase
95dl = (0.01008, 0.01003, 0.00990)
96
97{-# INLINE initPhase2 #-}
98initPhase2 :: (Phase, Phase)
99initPhase2 =
100   ((0,0.7,0.1), (0.3,0.4,0.6))
101
102
103size :: Int
104size = 10000000
105
106
107mainMonolithic1Composed :: IO ()
108mainMonolithic1Composed =
109   writeFile "speed.f32"
110      (fst $ unfoldrN size
111          (let (f0,f1,f2) = dl
112           in  generator0Freq f0 `mix`
113               generator0Freq f1 `mix`
114               generator0Freq f2)
115          (let (p0,p1,p2) = fst initPhase2
116           in  ((p0,p1),p2)))
117
118mainMonolithic1Generator :: IO ()
119mainMonolithic1Generator =
120   writeFile "speed.f32"
121      (runGeneratorMonolithic size
122          (let (f0,f1,f2) = dl
123               (p0,p1,p2) = fst initPhase2
124           in  generator0Gen f0 p0 `mixGen`
125               generator0Gen f1 p1 `mixGen`
126               generator0Gen f2 p2))
127
128main :: IO ()
129main = do args <- getArgs
130          case args of
131              ["1"] -> mainMonolithic1Generator
132              ["2"] -> mainMonolithic1Composed
133              _ -> error "Huh?"
134
135empty :: (Storable a) => Vector a
136empty = unsafeCreate 0 $ const $ return ()
137{-# NOINLINE empty #-}
138
139null :: Vector a -> Bool
140null (SV _ _ l) = assert (l >= 0) $ l <= 0
141{-# INLINE null #-}
142
143unfoldrN :: (Storable b) => Int -> (a -> Maybe (b, a)) -> a -> (Vector b, Maybe a)
144unfoldrN n f x0 =
145   if n <= 0
146     then (empty, Just x0)
147     else unsafePerformIO $ createAndTrim' n $ \p -> go p n x0
148       where
149          go = arguments2 $ \p i -> \x ->
150             if i == 0
151               then return (0, n-i, Just x)
152               else
153                 case f x of
154                   Nothing     -> return (0, n-i, Nothing)
155                   Just (w,x') -> do poke p w
156                                     go (incPtr p) (i-1) x'
157{-# INLINE unfoldrN #-}
158
159hPut :: (Storable a) => Handle -> Vector a -> IO ()
160hPut h v =
161   when (not (null v)) $
162      withStartPtr v $ \ ptrS l ->
163         let ptrE = advancePtr ptrS l
164         in  hPutBuf h ptrS (minusPtr ptrE ptrS)
165
166writeFile :: (Storable a) => FilePath -> Vector a -> IO ()
167writeFile f txt =
168   bracket (openBinaryFile f WriteMode) hClose
169      (\h -> hPut h txt)
170
171data Vector a = SV {-# UNPACK #-} !(ForeignPtr a)
172                   {-# UNPACK #-} !Int                -- offset
173                   {-# UNPACK #-} !Int                -- length
174
175withStartPtr :: Storable a => Vector a -> (Ptr a -> Int -> IO b) -> IO b
176withStartPtr (SV x s l) f =
177   withForeignPtr x $ \p -> f (p `advancePtr` s) l
178{-# INLINE withStartPtr #-}
179
180incPtr :: (Storable a) => Ptr a -> Ptr a
181incPtr v = advancePtr v 1
182{-# INLINE incPtr #-}
183
184unsafeCreate :: (Storable a) => Int -> (Ptr a -> IO ()) -> Vector a
185unsafeCreate l f = unsafePerformIO (create l f)
186{-# INLINE unsafeCreate #-}
187
188create :: (Storable a) => Int -> (Ptr a -> IO ()) -> IO (Vector a)
189create l f = do
190    fp <- mallocForeignPtrArray l
191    withForeignPtr fp $ \p -> f p
192    return $! SV fp 0 l
193
194createAndTrim' :: (Storable a) => Int
195                               -> (Ptr a -> IO (Int, Int, b))
196                               -> IO (Vector a, b)
197createAndTrim' l f = do
198    fp <- mallocForeignPtrArray l
199    withForeignPtr fp $ \p -> do
200        (off, l', res) <- f p
201        if assert (l' <= l) $ l' >= l
202            then return $! (SV fp 0 l, res)
203            else do ps <- create l' $ \p' -> copyArray p' (p `advancePtr` off) l'
204                    return $! (ps, res)
205
206{-# INLINE arguments2 #-}
207arguments2 :: (a -> b -> x) -> a -> b -> x
208arguments2 f a b = (f $! a) $! b
209
210{-# INLINE mallocForeignPtrArray #-}
211mallocForeignPtrArray :: Storable a => Int -> IO (F.ForeignPtr a)
212mallocForeignPtrArray = F.mallocForeignPtrArray