Does Haskell provides a way to evaluate IO monad immediately? - haskell

Question

I am currently making a ray tracing program with Haskell. As I am a very beginner of Haskell, I don't understand the evaluation strategy of IO monad clearly.
The problem is the memory usage of a long list of "IO a", which is "IO Vec" in my code.
Each element of the list is computed by a recursive function that compute IO Vec which represents the color for a pixel. Therefore, the length of the list is equals to width x height.
In addition, I take multiple samples for a pixels. As a whole, the function radiance to compute pixel value is called width x height x samples times.
First I was implemented this program simply by using list comprehension. The code is like,
main = do
...
let ray = (compute ray for every pair of [0..w-1], [0..h-1]
pixels <- sequence [ (sumOfRadiance scene ray samples) | ray <- rays]
In my understanding, as pixels is not used before it is written to a file, Haskell stores some data for function call inside pixels which is an array of IO Vec. Finally, memory consumption increases by calling recursive function radiance to compute pixel values.
If I change the program to evaluate the pixel value one by one using unsafePerformIO can prevent this strange use of memory space.
main = do
...
let ray = (compute ray for every pair of [0..w-1], [0..h-1]
let pixels = [ (unsafePerformIO (sumOfRadiance scene ray samples)) | ray <- rays]
I know unsafePerformIO is a bad solution, so I'd like to know if Haskell provides another way to evaluate inside of IO monad immediately. The following is the whole of my code (Sorry, it's a bit long...)
Thank you for your help.
-- Small path tracing with Haskell
import System.Environment
import System.Random.Mersenne
import System.IO.Unsafe
import Control.Monad
import Codec.Picture
import Data.Time
import qualified Data.Word as W
import qualified Data.Vector.Storable as V
-- Parameters
eps :: Double
eps = 1.0e-4
inf :: Double
inf = 1.0e20
nc :: Double
nc = 1.0
nt :: Double
nt = 1.5
-- Vec
data Vec = Vec (Double, Double, Double) deriving (Show)
instance (Num Vec) where
(Vec (x, y, z)) + (Vec (a, b, c)) = Vec (x + a, y + b, z + c)
(Vec (x, y, z)) - (Vec (a, b, c)) = Vec (x - a, y - b, z - c)
(Vec (x, y, z)) * (Vec (a, b, c)) = Vec (x * a, y * b, z * c)
abs = undefined
signum = undefined
fromInteger x = Vec (dx, dx, dx) where dx = fromIntegral x
x :: Vec -> Double
x (Vec (x, _, _)) = x
y :: Vec -> Double
y (Vec (_, y, _)) = y
z :: Vec -> Double
z (Vec (_, _, z)) = z
mul :: Vec -> Double -> Vec
mul (Vec (x, y, z)) s = Vec (x * s, y * s, z * s)
dot :: Vec -> Vec -> Double
dot (Vec (x, y, z)) (Vec (a, b, c)) = x * a + y * b + z * c
norm :: Vec -> Vec
norm (Vec (x, y, z)) = Vec (x * invnrm, y * invnrm, z * invnrm)
where invnrm = 1 / sqrt (x * x + y * y + z * z)
cross :: Vec -> Vec -> Vec
cross (Vec (x, y, z)) (Vec (a, b, c)) = Vec (y * c - b * z, z * a - c * x, x * b - a * y)
-- Ray
data Ray = Ray (Vec, Vec) deriving (Show)
org :: Ray -> Vec
org (Ray (org, _)) = org
dir :: Ray -> Vec
dir (Ray (_, dir)) = dir
-- Material
data Refl = Diff
| Spec
| Refr
deriving Show
-- Sphere
data Sphere = Sphere (Double, Vec, Vec, Vec, Refl) deriving (Show)
rad :: Sphere -> Double
rad (Sphere (rad, _, _, _, _ )) = rad
pos :: Sphere -> Vec
pos (Sphere (_ , p, _, _, _ )) = p
emit :: Sphere -> Vec
emit (Sphere (_ , _, e, _, _ )) = e
col :: Sphere -> Vec
col (Sphere (_ , _, _, c, _ )) = c
refl :: Sphere -> Refl
refl (Sphere (_ , _, _, _, refl)) = refl
intersect :: Sphere -> Ray -> Double
intersect sp ray =
let op = (pos sp) - (org ray)
b = op `dot` (dir ray)
det = b * b - (op `dot` op) + ((rad sp) ** 2)
in
if det < 0.0
then inf
else
let sqdet = sqrt det
t1 = b - sqdet
t2 = b + sqdet
in ansCheck t1 t2
where ansCheck t1 t2
| t1 > eps = t1
| t2 > eps = t2
| otherwise = inf
-- Scene
type Scene = [Sphere]
sph :: Scene
sph = [ Sphere (1e5, Vec ( 1e5+1, 40.8, 81.6), Vec (0.0, 0.0, 0.0), Vec (0.75, 0.25, 0.25), Diff) -- Left
, Sphere (1e5, Vec (-1e5+99, 40.8, 81.6), Vec (0.0, 0.0, 0.0), Vec (0.25, 0.25, 0.75), Diff) -- Right
, Sphere (1e5, Vec (50.0, 40.8, 1e5), Vec (0.0, 0.0, 0.0), Vec (0.75, 0.75, 0.75), Diff) -- Back
, Sphere (1e5, Vec (50.0, 40.8, -1e5+170), Vec (0.0, 0.0, 0.0), Vec (0.0, 0.0, 0.0), Diff) -- Front
, Sphere (1e5, Vec (50, 1e5, 81.6), Vec (0.0, 0.0, 0.0), Vec (0.75, 0.75, 0.75), Diff) -- Bottom
, Sphere (1e5, Vec (50,-1e5+81.6,81.6), Vec (0.0, 0.0, 0.0), Vec (0.75, 0.75, 0.75), Diff) -- Top
, Sphere (16.5, Vec (27, 16.5, 47), Vec (0.0, 0.0, 0.0), Vec (1,1,1) `mul` 0.999, Spec) -- Mirror
, Sphere (16.5, Vec (73, 16.5, 78), Vec (0.0, 0.0, 0.0), Vec (1,1,1) `mul` 0.999, Refr) -- Glass
, Sphere (600, Vec (50, 681.6 - 0.27, 81.6), Vec (12, 12, 12), Vec (0, 0, 0), Diff) ] -- Light
-- Utility functions
clamp :: Double -> Double
clamp = (max 0.0) . (min 1.0)
isectWithScene :: Scene -> Ray -> (Double, Int)
isectWithScene scene ray = foldr1 (min) $ zip [ intersect sph ray | sph <- scene ] [0..]
nextDouble :: IO Double
nextDouble = randomIO
lambert :: Vec -> Double -> Double -> (Vec, Double)
lambert n r1 r2 =
let th = 2.0 * pi * r1
r2s = sqrt r2
w = n
u = norm $ (if (abs (x w)) > eps then Vec (0, 1, 0) else Vec (1, 0, 0)) `cross` w
v = w `cross` u
uu = u `mul` ((cos th) * r2s)
vv = v `mul` ((sin th) * r2s)
ww = w `mul` (sqrt (1.0 - r2))
rdir = norm (uu + vv + ww)
in (rdir, 1)
reflect :: Vec -> Vec -> (Vec, Double)
reflect v n =
let rdir = v - (n `mul` (2.0 * n `dot` v))
in (rdir, 1)
refract :: Vec -> Vec -> Vec -> Double -> (Vec, Double)
refract v n orn rr =
let (rdir, _) = reflect v orn
into = (n `dot` orn) > 0
nnt = if into then (nc / nt) else (nt / nc)
ddn = v `dot` orn
cos2t = 1.0 - nnt * nnt * (1.0 - ddn * ddn)
in
if cos2t < 0.0
then (rdir, 1.0)
else
let tdir = norm $ ((v `mul` nnt) -) $ n `mul` ((if into then 1 else -1) * (ddn * nnt + (sqrt cos2t)))
a = nt - nc
b = nt + nc
r0 = (a * a) / (b * b)
c = 1.0 - (if into then -ddn else (tdir `dot` n))
re = r0 + (1 - r0) * (c ** 5)
tr = 1.0 - re
pp = 0.25 + 0.5 * re
in
if rr < pp
then (rdir, (pp / re))
else (tdir, ((1.0 - pp) / tr))
radiance :: Scene -> Ray -> Int -> IO Vec
radiance scene ray depth = do
let (t, i) = (isectWithScene scene ray)
if inf <= t
then return (Vec (0, 0, 0))
else do
r0 <- nextDouble
r1 <- nextDouble
r2 <- nextDouble
let obj = (scene !! i)
let c = col obj
let prob = (max (x c) (max (y c) (z c)))
if depth >= 5 && r0 >= prob
then return (emit obj)
else do
let rlt = if depth < 5 then 1 else prob
let f = (col obj)
let d = (dir ray)
let x = (org ray) + (d `mul` t)
let n = norm $ x - (pos obj)
let orn = if (d `dot` n) < 0.0 then n else (-n)
let (ndir, pdf) = case (refl obj) of
Diff -> (lambert orn r1 r2)
Spec -> (reflect d orn)
Refr -> (refract d n orn r1)
nextRad <- (radiance scene (Ray (x, ndir)) (succ depth))
return $ ((emit obj) + ((f * nextRad) `mul` (1.0 / (rlt * pdf))))
toByte :: Double -> W.Word8
toByte x = truncate (((clamp x) ** (1.0 / 2.2)) * 255.0) :: W.Word8
accumulateRadiance :: Scene -> Ray -> Int -> Int -> IO Vec
accumulateRadiance scene ray d m = do
let rays = take m $ repeat ray
pixels <- sequence [radiance scene r 0 | r <- rays]
return $ (foldr1 (+) pixels) `mul` (1 / fromIntegral m)
main :: IO ()
main = do
args <- getArgs
let argc = length args
let w = if argc >= 1 then (read (args !! 0)) else 400 :: Int
let h = if argc >= 2 then (read (args !! 1)) else 300 :: Int
let spp = if argc >= 3 then (read (args !! 2)) else 4 :: Int
startTime <- getCurrentTime
putStrLn "-- Smallpt.hs --"
putStrLn $ " width = " ++ (show w)
putStrLn $ " height = " ++ (show h)
putStrLn $ " spp = " ++ (show spp)
let dw = fromIntegral w :: Double
let dh = fromIntegral h :: Double
let cam = Ray (Vec (50, 52, 295.6), (norm $ Vec (0, -0.042612, -1)));
let cx = Vec (dw * 0.5135 / dh, 0.0, 0.0)
let cy = (norm $ cx `cross` (dir cam)) `mul` 0.5135
let dirs = [ norm $ (dir cam) + (cy `mul` (y / dh - 0.5)) + (cx `mul` (x / dw - 0.5)) | y <- [dh-1,dh-2..0], x <- [0..dw-1] ]
let rays = [ Ray ((org cam) + (d `mul` 140.0), (norm d)) | d <- dirs ]
let pixels = [ (unsafePerformIO (accumulateRadiance sph r 0 spp)) | r <- rays ]
let pixelData = map toByte $! pixels `seq` (foldr (\col lst -> [(x col), (y col), (z col)] ++ lst) [] pixels)
let pixelBytes = V.fromList pixelData :: V.Vector W.Word8
let img = Image { imageHeight = h, imageWidth = w, imageData = pixelBytes } :: Image PixelRGB8
writePng "image.png" img
endTime <- getCurrentTime
print $ diffUTCTime endTime startTime

Answer 1

First, I think there is an error. When you talk about going from
pixels <- sequence [ (sumOfRadiance scene ray samples) | ray <- rays]
to
pixels <- sequence [ (unsafePerformIO (sumOfRadiance scene ray samples)) | ray <- rays]
that doesn't make sense. The types shouldn't match up -- sequence only makes sense if you are combining a bunch of things of type m a. It would be correct to do
let pixels = [ unsafePerformIO (sumOfRadiance scene ray samples) | ray <- rays ]
I will somewhat cavalierly assume that that is what you did and you simply made a mistake when entering your question.
If this is the case, then what you are actually looking for is a way to execute IO actions more lazily, not more immediately. The sequence call forces all the actions to be run right then, whereas the unsafePerformIO version simply creates a list of un-run actions (and indeed the list itself is generated lazily so it doesn't exist all at once), and the actions are run individually as their results are needed.
It appears that the reason you need IO is to generate random numbers. Randomness can be kind of a pain -- usually MonadRandom does the job, but it still creates a sequential dependence between actions and may still not be lazy enough (I'd give it a try -- if you use it you get reproducibility -- the same seed gives the same results, even after refactorings that respect the monad laws).
If MonadRandom doesn't work and you need to generate random numbers in a more on-demand way, the way would be to make your own randomness monad which does the same thing as your unsafePerformIO solution, but in a way that is properly encapsulated. I'm going to show you the way I consider to be the Haskell Way To Cheat. First, a lovely pure implementation sketch:
-- A seed tells you how to generate random numbers
data Seed = ...
splitSeed :: Seed -> (Seed, Seed)
random :: Seed -> Double
-- A Cloud is a probability distribution of a's, or an a which
-- depends on a random seed. This monad is just as lazy as a
-- pure computation.
newtype Cloud a = Cloud { runCloud :: Seed -> a }
deriving (Functor)
instance Monad Cloud where
return = Cloud . const
m >>= f = Cloud $ \seed ->
let (seed1, seed2) = splitSeed seed in
runCloud (f (runCloud m seed1)) seed2
(I think I got that right. The point is that at every bind you split the seed in two and pass one to the left and the other to the right.)
Now this is a perfectly pure implementation of randomness... with a couple catches. (1) there is no non-trivial splitSeed which will strictly respect the monad laws, and (2) even if we allow the laws to be broken, random number generators based on splitting can be pretty slow. But if we give up determinism, if all we care about is that we get a good sampling from the distribution rather than the exact same result, then we don't need to strictly respect the monad laws. And at that point we cheat and pretend there is a suitable Seed type:
data Seed = Seed
splitSeed Seed = (Seed, Seed)
-- Always NOINLINE functions with unsafePerformIO to keep the
-- optimizer from messing with you.
{-# NOINLINE random #-}
random Seed = unsafePerformIO randomIO
We should hide this inside a module to keep the abstraction barrier clear. Cloud and runCloud should not be exposed since they allow us to violate purity; expose only
runCloudIO :: Cloud a -> IO a
runCloudIO = return . runCloud
which doesn't technically need IO, but communicates that this will not be deterministic. Then you can build up whatever you need as a value in the Cloud monad, and run it once in your main program.
You might ask why we have a Seed type at all if it doesn't have any information. Well, I think splitSeed is just a nod to purity and isn't actually doing anything -- you could remove it -- but we need Cloud to be a function type so that the implicit caching of laziness doesn't break our semantics. Otherwise
let foo = random in liftM2 (,) foo foo
would always return a pair with two identical components, since the random value was really associated with foo. I am not sure about these things since at this point we are at war with the optimizer, it takes some experimentation.
Happy cheating. :-)