Initial

6 files changed, 329 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..0b605e0
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,4 @@
+dist-newstyle/
+accelerate-llvm/
+out.data
+out.png
diff --git a/cabal.project b/cabal.project
new file mode 100644
index 0000000..247bfa3
--- /dev/null
+++ b/cabal.project
@@ -0,0 +1,28 @@
+packages:
+  .
+  accelerate-llvm/accelerate-llvm
+  accelerate-llvm/accelerate-llvm-native
+
+allow-newer:
+  accelerate-llvm:llvm-hs,
+  accelerate-llvm:llvm-hs-pure,
+  accelerate-llvm-native:llvm-hs,
+  accelerate-llvm-native:llvm-hs-pure
+
+source-repository-package
+  type: git
+  location: https://github.com/llvm-hs/llvm-hs
+  tag: llvm-12
+  subdir: llvm-hs llvm-hs-pure
+
+-- source-repository-package
+--   type: git
+--   location: https://github.com/AccelerateHS/accelerate-llvm
+--   tag: master
+--   subdir: accelerate-llvm accelerate-llvm-native accelerate-llvm-ptx
+
+source-repository-package
+  type: git
+  location: https://github.com/AccelerateHS/accelerate
+  tag: master
+  post-checkout-command: bash -c "! git submodule update --init --recursive"
diff --git a/compute/Compute.hs b/compute/Compute.hs
new file mode 100644
index 0000000..fde31cf
--- /dev/null
+++ b/compute/Compute.hs
@@ -0,0 +1,148 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE ViewPatterns #-}
+module Main where
+
+import qualified Data.Array.Accelerate as A
+-- import qualified Data.Array.Accelerate.Interpreter as Interpreter
+import qualified Data.Array.Accelerate.LLVM.Native as CPU
+import qualified Data.ByteString.Lazy as BSL
+import Data.Function ((&))
+import qualified Data.Serialize as Ser
+import Data.Word
+import qualified Data.Vector.Storable as VS
+import Text.Read (readMaybe)
+import System.Environment (getArgs)
+import System.Exit (exitFailure)
+
+import MandelHSlib
+
+
+mandelstep :: (A.Ord a, A.Num a, A.Elt a) => A.Exp (a, a) -> A.Exp (a, a) -> A.Exp (a, a)
+mandelstep (A.T2 cx cy) (A.T2 x y) = A.T2 (x * x - y * y + cx) (2 * x * y + cy)
+
+mandeliter :: (A.Ord a, A.Num a, A.Elt a) => A.Exp (a, a) -> Int -> A.Exp Int
+mandeliter (A.T2 cx cy) maxiter =
+  let A.T3 nres _ _ =
+        A.while (\(A.T3 n x y) -> n A.< A.constant maxiter A.&& x * x + y * y A.< 4)
+                (\(A.T3 n x y) ->
+                    let A.T2 x' y' = mandelstep (A.T2 cx cy) (A.T2 x y)
+                    in A.T3 (n+1) x' y')
+                (A.T3 0 cx cy)
+  in nres
+
+pixel :: (A.Num a, Num a, Fractional a, A.Floating a, A.ToFloating Int a, A.Elt a)
+      => (Int, Int) -> A.Exp (a, a) -> A.Exp a -> A.Exp A.DIM2 -> A.Exp (a, a)
+pixel (w, h) (A.T2 cx cy) cw (A.I2 iy ix) =
+  let ch = cw * A.constant (fromIntegral h / fromIntegral w)
+      ltx = cx - cw / 2
+      lty = cy - ch / 2
+  in A.T2 (ltx + cw / fromIntegral (w - 1) * A.toFloating ix)
+          (lty + ch / fromIntegral (h - 1) * A.toFloating iy)
+
+invpixel :: (A.Num a, Num a, Fractional a, A.RealFrac a, A.ToFloating Int a, A.Elt a)
+         => (Int, Int) -> A.Exp (a, a) -> A.Exp a -> A.Exp (a, a) -> A.Exp (Maybe A.DIM2)
+invpixel (w, h) (A.T2 cx cy) cw (A.T2 x y) =
+  let ch = cw * A.constant (fromIntegral h / fromIntegral w)
+      ltx = cx - cw / 2
+      lty = cy - ch / 2
+      ix = A.round $ (x - ltx) / cw * fromIntegral (w - 1)
+      iy = A.round $ (y - lty) / ch * fromIntegral (h - 1)
+  in A.cond (0 A.<= ix A.&& ix A.< A.constant w A.&&
+             0 A.<= iy A.&& iy A.< A.constant h)
+       (A.Just_ (A.I2 iy ix))
+       A.Nothing_
+
+mandelbrot :: (A.Num a, Num a, A.Ord a, Fractional a, A.Floating a, A.ToFloating Int a, A.Elt a)
+           => (Int, Int) -> Int -> A.Acc (A.Scalar (a, a)) -> A.Acc (A.Scalar a)
+           -> A.Acc (A.Matrix Int)
+mandelbrot size@(w, h) maxiter (A.the -> cpos) (A.the -> cw) =
+  A.generate (A.I2 (A.constant h) (A.constant w)) $ \idx ->
+    mandeliter (pixel size cpos cw idx) maxiter
+
+mandelSet :: (A.Num a, Num a, A.Ord a, Fractional a, A.Floating a, A.ToFloating Int a, A.Elt a)
+          => (Int, Int) -> Int -> A.Acc (A.Scalar (a, a)) -> A.Acc (A.Scalar a)
+          -> A.Acc (A.Matrix Bool)
+mandelSet size maxiter cpos cw =
+  A.map (A.== A.constant maxiter) (mandelbrot size maxiter cpos cw)
+
+image :: (A.Num a, Num a, A.Ord a, Fractional a, A.RealFrac a, A.Floating a, A.ToFloating Int a, A.Elt a)
+      => (Int, Int) -> Int -> A.Acc (A.Scalar (a, a)) -> A.Acc (A.Scalar a)
+      -> A.Acc (A.Matrix a)
+image size@(w, h) maxiter acpos@(A.the -> cpos) acw@(A.the -> cw) =
+  let mbrot = mandelbrot (w, h) maxiter acpos acw
+      mset = A.map (A.== A.constant maxiter) mbrot
+      isize = A.I2 (A.constant h) (A.constant w)
+      pixels = A.generate isize (pixel size cpos cw)
+      A.T2 _ res =
+        A.awhile (\(A.T2 (A.the -> iter) _) -> A.unit (iter A.< A.constant maxiter))
+                 (\(A.T2 iter current) ->
+                   let current' =
+                         A.zipWith (\c (A.T2 z n) ->
+                                      let z' = mandelstep c z
+                                          n' = invpixel size cpos cw z' & A.match \case
+                                                 A.Nothing_ -> n
+                                                 A.Just_ z'pix ->
+                                                   A.cond (mset A.! z'pix) (n+1) n
+                                      in A.T2 z' n')
+                                   pixels current
+                   in A.T2 (A.map (+1) iter) current')
+                 (A.T2 (A.unit 0)
+                       (A.map (\z -> A.T2 z (0 :: A.Exp Int)) pixels))
+      counts = A.map A.snd res
+  in A.zipWith (/) (A.map A.toFloating counts) (A.map A.toFloating mbrot)
+
+
+colorscheme :: (Ord a, Floating a, RealFrac a) => a -> (Word8, Word8, Word8)
+colorscheme fraction
+  -- | fraction == 1.0 = (0, 0, 0)
+  | otherwise =
+      let x = max 0 fraction ** 0.3
+          bg = 0.2 * (1 - curve x (-0.1) 0.2)
+      in (tow8 $ bg + curve x 0 0.4 - 0.8 * curve x 0.6 1.0
+         ,tow8 $ bg + 0.8 * curve x 0.3 0.7
+         ,tow8 $ bg + curve x 0.6 1.0)
+  where
+    tow8 x = round (max 0 (min 255 (x * 255)))
+    curve x start end
+      | x <= start = 0
+      | x >= end = 1
+      | otherwise = sin (pi/(end-start) * (x - start) - pi/2) / 2 + 0.5
+
+computeFractal :: (Ord a, Floating a, RealFrac a, A.Elt a, A.RealFrac a, A.Floating a, A.ToFloating Int a, VS.Storable a)
+               => (Int, Int) -> Int -> Fractal a
+computeFractal (w, h) maxiter =
+  -- let arr = CPU.runN (\cpos cw ->
+  --              A.map (\n -> A.toFloating n / A.toFloating (A.constant maxiter))
+  --                    (mandelbrot (w, h) maxiter cpos cw))
+  --               (A.fromList A.Z [(-0.5, 0.0)])
+  --               (A.fromList A.Z [3.5])
+  let arr = CPU.runN (\cpos cw -> image (w, h) maxiter cpos cw)
+                (A.fromList A.Z [(-0.5, 0.0)])
+                (A.fromList A.Z [3.5])
+  in Fractal (w, h) maxiter (VS.fromList (A.toList arr))
+
+usage :: String
+usage =
+  "Usage: mandelhs-compute <outfile.data> <width> <height> <maxiter>\n"
+
+main :: IO ()
+main = do
+  -- let (w, h) = (4 * 2000, 4 * 1500)
+  --     maxiter = 1024
+
+  args <- getArgs
+  case args of
+    [outfile, wids, heis, maxiters] ->
+      case (readMaybe wids, readMaybe heis, readMaybe maxiters) of
+        (Just w, Just h, Just maxiter) ->
+          BSL.writeFile outfile $
+            Ser.runPutLazy (Ser.put (computeFractal @Double (w, h) maxiter))
+        _ -> putStr usage >> exitFailure
+
+    _ -> putStr usage >> exitFailure
diff --git a/lib/MandelHSlib.hs b/lib/MandelHSlib.hs
new file mode 100644
index 0000000..d9409f7
--- /dev/null
+++ b/lib/MandelHSlib.hs
@@ -0,0 +1,20 @@
+module MandelHSlib where
+
+import qualified Data.Serialize as Ser
+import qualified Data.Vector.Storable as VS
+
+
+data Fractal a = Fractal (Int, Int) Int (VS.Vector a)
+
+instance (VS.Storable a, Ser.Serialize a) => Ser.Serialize (Fractal a) where
+  put (Fractal size maxiter v) = do
+    Ser.put size
+    Ser.put maxiter
+    Ser.putWord64be (fromIntegral (VS.length v))
+    VS.mapM_ Ser.put v
+  get = do
+    size <- Ser.get
+    maxiter <- Ser.get
+    len <- fromIntegral <$> Ser.getWord64be
+    v <- VS.replicateM len Ser.get
+    return (Fractal size maxiter v)
diff --git a/mandelhs.cabal b/mandelhs.cabal
new file mode 100644
index 0000000..8e92db9
--- /dev/null
+++ b/mandelhs.cabal
@@ -0,0 +1,62 @@
+cabal-version:       2.0
+name:                mandelhs
+synopsis:            Mandelbrots
+version:             0.1.0.0
+license:             MIT
+author:              Tom Smeding
+maintainer:          tom@tomsmeding.com
+build-type:          Simple
+
+library mandelhs-lib
+  exposed-modules:
+    MandelHSlib
+  build-depends:
+    base >= 4.13 && < 4.15,
+    -- accelerate-llvm-native already depends on cereal
+    cereal,
+    vector
+  hs-source-dirs:
+    lib
+  default-language:
+    Haskell2010
+  ghc-options:
+    -Wall -O2
+
+executable mandelhs-compute
+  main-is:
+    Compute.hs
+  other-modules:
+  build-depends:
+    mandelhs-lib,
+    base >= 4.13 && < 4.15,
+    accelerate ^>= 1.3.0.0,
+    accelerate-llvm-native ^>= 1.3.0.0,
+    bytestring,
+    -- accelerate-llvm-native already depends on cereal
+    cereal,
+    vector
+  hs-source-dirs:
+    compute
+  default-language:
+    Haskell2010
+  ghc-options:
+    -Wall -O2 -threaded
+
+executable mandelhs-render
+  main-is:
+    Render.hs
+  other-modules:
+  build-depends:
+    mandelhs-lib,
+    base >= 4.13 && < 4.15,
+    bytestring,
+    -- accelerate-llvm-native already depends on cereal
+    cereal,
+    JuicyPixels ^>= 3.3.8,
+    vector
+  hs-source-dirs:
+    render
+  default-language:
+    Haskell2010
+  ghc-options:
+    -Wall -O2 -threaded
diff --git a/render/Render.hs b/render/Render.hs
new file mode 100644
index 0000000..9126032
--- /dev/null
+++ b/render/Render.hs
@@ -0,0 +1,67 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE ViewPatterns #-}
+module Main where
+
+import qualified Codec.Picture.Png as PNG
+import qualified Codec.Picture.Types as Pix
+import qualified Data.ByteString.Lazy as BSL
+import qualified Data.Serialize as Ser
+import Data.Word
+import qualified Data.Vector.Storable as VS
+import System.Environment (getArgs)
+import System.Exit (die, exitFailure)
+
+import MandelHSlib
+
+
+colorscheme :: (Ord a, Floating a, RealFrac a) => a -> (Word8, Word8, Word8)
+colorscheme fraction =
+  let x = max 0 fraction ** 0.3
+      bg = 0.2 * (1 - curve x (-0.1) 0.2)
+  in (tow8 $ bg + curve x 0 0.4 - 0.8 * curve x 0.6 1.0
+     ,tow8 $ bg + 0.8 * curve x 0.3 0.7
+     ,tow8 $ bg + curve x 0.6 1.0)
+  where
+    tow8 x = round (max 0 (min 255 (x * 255)))
+    curve x start end
+      | x <= start = 0
+      | x >= end = 1
+      | otherwise = sin (pi/(end-start) * (x - start) - pi/2) / 2 + 0.5
+
+renderFractal :: (Ord a, Floating a, RealFrac a, VS.Storable a) => Fractal a -> Pix.Image Pix.PixelRGB8
+renderFractal (Fractal (w, h) _ arr) =
+  Pix.Image w h $ VS.fromList
+    [val
+    | fraction <- VS.toList arr
+    , let (r, g, b) = colorscheme fraction
+    , val <- [r, g, b]]
+
+usage :: String
+usage =
+  "Usage: mandelhs-render <infile.data> <outfile.png>\n"
+
+main :: IO ()
+main = do
+  -- let (w, h) = (4 * 2000, 4 * 1500)
+  --     maxiter = 1024
+
+  args <- getArgs
+  case args of
+    [infile, outfile] -> do
+      datafile <- BSL.readFile infile
+      Fractal (w, h) _ arr <- case Ser.runGetLazy Ser.get datafile of
+        Left err -> die err
+        Right res -> return res
+      BSL.writeFile outfile $
+        PNG.encodePng $
+          Pix.Image @Pix.PixelRGB8 w h $ VS.fromList
+            [val
+            | fraction <- VS.toList arr
+            , let (r, g, b) = colorscheme @Double fraction
+            , val <- [r, g, b]]
+
+    _ -> putStr usage >> exitFailure