packages feed

accelerate-cublas-0.1: benchmark/NewtonInverse.hs

{-# LANGUAGE FlexibleContexts #-}
module Main where

import qualified Data.Array.Accelerate.CUBLAS.Level3.Batched as Batched

import qualified Data.Array.Accelerate.LinearAlgebra as ALinAlg
import qualified Data.Array.Accelerate.Utility.Loop as Loop
import qualified Data.Array.Accelerate.CUDA as CUDA
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate (All(All), Z(Z), (:.)((:.)))

import qualified Control.Concurrent.PooledIO.Independent as Pooled

import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import qualified Numeric.Container as Container
import qualified Numeric.LinearAlgebra.Algorithms as HMLinAlg

import Numeric.Container (Container, (<>))
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import qualified System.Random as Rnd
import System.TimeIt (timeIt)

import Text.Printf (printf)

import qualified Data.List.HT as ListHT
import Data.Function.HT (nest)
import Data.Tuple.HT (mapPair)



newtonInverseStep ::
   (Num a, Container Vector a, Container.Product a) =>
   Matrix a -> Matrix a -> Matrix a
newtonInverseStep a x =
   Container.sub
      (Container.scale 2 x)
      (x <> a <> x)

newtonInverse ::
   (Num a, Container Vector a, Container.Product a) =>
   Int -> Matrix a -> Matrix a -> Matrix a
newtonInverse count start a =
   nest count (newtonInverseStep a) start


newtonInverseCUBLASStep, newtonInverseCUBLASStepMul ::
   (A.Shape ix, A.Slice ix, Eq ix, Batched.Element a, A.IsNum a, A.Elt a) =>
   Batched.Handle ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a
newtonInverseCUBLASStep h a x =
   Batched.mac h (-1) x (Batched.mul h 1 a x) 2 x

newtonInverseCUBLASStepMul h a x =
   A.zipWith (-) (A.map (2*) x) $
   Batched.mul h 1 x $ Batched.mul h 1 a x

newtonInverseCUBLAS ::
   (A.Shape ix, A.Slice ix, Eq ix, Batched.Element a, A.IsNum a, A.Elt a) =>
   Batched.Handle ->
   A.Exp Int ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a
newtonInverseCUBLAS h n seed a =
   Loop.nest n (newtonInverseCUBLASStep h a) seed


randomMatrixInv :: Int -> (Matrix Double, Matrix Double)
randomMatrixInv size =
   let x =
          Matrix.fromLists $ take size $ ListHT.sliceVertical size $
          Rnd.randomRs (-1,1::Double) $ Rnd.mkStdGen 42
   in  (x, HMLinAlg.inv x)


parallel :: [a] -> (Int -> a -> IO ()) -> IO ()
parallel xs f = Pooled.run $ zipWith f [0 ..] xs

disturbedMatrices ::
   (Container Vector a) =>
   Matrix a -> [a] -> [Matrix a]
disturbedMatrices x yelems =
   let size = Matrix.rows x
   in  map (Container.add x . Matrix.fromLists) $
       ListHT.sliceVertical size $
       ListHT.sliceVertical size $
       yelems

mainHMatrixDirect ::
   (Show a, Container Vector a, HMLinAlg.Field a) =>
   String ->
   Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainHMatrixDirect typ numberOfMatrices (x, _xinv) yelems = do
   let yinvs = map HMLinAlg.inv $ disturbedMatrices x yelems
   putStrLn $ "hmatrix-direct-" ++ typ
   timeIt $ parallel (take numberOfMatrices yinvs) $ \ n y ->
      writeFile (printf "/tmp/hmatrix-direct-%s%03d.txt" typ n) $ show y

{-
das wird durch mehrere Kerne gar nicht schneller,
vielleicht st�rt das Serialisieren der Zufallszahlen?
-}
mainHMatrix ::
   (Show a, Container Vector a, Container.Product a) =>
   String ->
   Int -> Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainHMatrix typ numberOfMatrices newtonIts (x, xinv) yelems = do
   let yinvs = map (newtonInverse newtonIts xinv) $ disturbedMatrices x yelems
   putStrLn $ "hmatrix-" ++ typ
   timeIt $ parallel (take numberOfMatrices yinvs) $ \ n y ->
      writeFile (printf "/tmp/hmatrix-%s%03d.txt" typ n) $ show y


mainCUDA ::
   (A.Elt a, A.IsNum a, Container.Element a) =>
   String ->
   Int -> Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainCUDA typ numberOfMatrices newtonIts (xm, xinvm) yelems = do
   let size = Matrix.rows xm
       matrixAccFromHM =
          A.fromList (Z :. size :. size) .
          Vector.toList . Matrix.flatten
       xarr = matrixAccFromHM xm
       xinvarr = matrixAccFromHM xinvm

   let ysarr =
          A.fromList (Z :. numberOfMatrices :. size :. size) yelems
       rep = A.replicate (A.lift $ Z :. numberOfMatrices :. All :. All)
       yinvs =
          CUDA.run1
             (\args ->
                case A.unlift args of
                   (x, xinv, ys) ->
                      ALinAlg.newtonInverse (A.constant newtonIts) (rep xinv) $
                      A.zipWith (+) ys (rep x))
             (xarr, xinvarr, ysarr)

   putStrLn $ "cuda-" ++ typ
   timeIt $ writeFile ("/tmp/cuda-"++typ++".txt") $ show yinvs


mainCUBLASDirect ::
   (Batched.Element a, Container.Element a, A.IsNum a, A.Elt a) =>
   String ->
   Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainCUBLASDirect typ numberOfMatrices (xm, _xinvm) yelems = do
   let size = Matrix.rows xm
       matrixAccFromHM =
          A.fromList (Z :. size :. size) .
          Vector.toList . Matrix.flatten
       xarr = matrixAccFromHM xm

   handle <- Batched.create
   let ysarr =
          A.fromList (Z :. numberOfMatrices :. size :. size) yelems
       rep = A.replicate (A.lift $ Z :. numberOfMatrices :. All :. All)
       yinvs =
          CUDA.run1
             (\args ->
                case A.unlift args of
                   (x, ys) ->
                      fst $ Batched.inv handle $ A.zipWith (+) ys (rep x))
             (xarr, ysarr)

   putStrLn $ "cublas-direct-" ++ typ
   timeIt $ writeFile ("/tmp/cublas-direct-"++typ++".txt") $ show yinvs


mainCUBLAS ::
   (Batched.Element a, Container.Element a, A.IsNum a, A.Elt a) =>
   String ->
   Int -> Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainCUBLAS typ numberOfMatrices newtonIts (xm, xinvm) yelems = do
   let size = Matrix.rows xm
       matrixAccFromHM =
          A.fromList (Z :. size :. size) .
          Vector.toList . Matrix.flatten
       xarr = matrixAccFromHM xm
       xinvarr = matrixAccFromHM xinvm

   handle <- Batched.create
   let ysarr =
          A.fromList (Z :. numberOfMatrices :. size :. size) yelems
       rep = A.replicate (A.lift $ Z :. numberOfMatrices :. All :. All)
       yinvs =
          CUDA.run1
             (\args ->
                case A.unlift args of
                   (x, xinv, ys) ->
                      newtonInverseCUBLAS handle (A.constant newtonIts) (rep xinv) $
                      A.zipWith (+) ys (rep x))
             (xarr, xinvarr, ysarr)

   putStrLn $ "cublas-" ++ typ
   timeIt $ writeFile ("/tmp/cublas-"++typ++".txt") $ show yinvs


main :: IO ()
main = do
   let n = 96
   let sz = 50
   let its = 20
   let xmsDouble = randomMatrixInv sz
       ysDouble = Rnd.randomRs (-0.01,0.01) $ Rnd.mkStdGen 23
   let xmsFloat =
          mapPair
             (Container.cmap realToFrac, Container.cmap realToFrac)
             xmsDouble
       ysFloat :: [Float]
       ysFloat = map realToFrac ysDouble
   mainHMatrixDirect "double" n xmsDouble ysDouble
   -- mainHMatrixDirect "float" n xmsFloat ysFloat
   mainHMatrix "double" n its xmsDouble ysDouble
   mainHMatrix "float" n its xmsFloat ysFloat
   mainCUBLASDirect "double" n xmsDouble ysDouble
   mainCUBLASDirect "float" n xmsFloat ysFloat
   mainCUBLAS "double" n its xmsDouble ysDouble
   mainCUBLAS "float" n its xmsFloat ysFloat
   mainCUDA "double" n its xmsDouble ysDouble
   mainCUDA "float" n its xmsFloat ysFloat