accelerate-cublas (empty) → 0.0
raw patch · 7 files changed
+745/−0 lines, 7 filesdep +acceleratedep +accelerate-arithmeticdep +accelerate-cublassetup-changed
Dependencies added: accelerate, accelerate-arithmetic, accelerate-cublas, accelerate-cuda, accelerate-io, accelerate-utility, base, cublas, cuda, random, utility-ht, vector
Files
- LICENSE +27/−0
- Setup.lhs +3/−0
- accelerate-cublas.cabal +65/−0
- example/Main.hs +109/−0
- src/Data/Array/Accelerate/CUBLAS/Level2/Batched.hs +42/−0
- src/Data/Array/Accelerate/CUBLAS/Level3/Batched.hs +266/−0
- src/Data/Array/Accelerate/CUBLAS/Level3/Batched/Foreign.hs +233/−0
+ LICENSE view
@@ -0,0 +1,27 @@+Copyright (c) Henning Thielemann 2014++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions+are met:+1. Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.+2. Redistributions in binary form must reproduce the above copyright+ notice, this list of conditions and the following disclaimer in the+ documentation and/or other materials provided with the distribution.+3. Neither the name of the author nor the names of his contributors+ may be used to endorse or promote products derived from this software+ without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE+ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS+OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)+HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT+LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY+OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF+SUCH DAMAGE.
+ Setup.lhs view
@@ -0,0 +1,3 @@+#! /usr/bin/env runhaskell+> import Distribution.Simple+> main = defaultMain
+ accelerate-cublas.cabal view
@@ -0,0 +1,65 @@+Name: accelerate-cublas+Version: 0.0+License: BSD3+License-File: LICENSE+Author: Henning Thielemann <haskell@henning-thielemann.de>+Maintainer: Henning Thielemann <haskell@henning-thielemann.de>+Homepage: http://code.haskell.org/~thielema/accelerate-cublas/+Category: Math+Synopsis: Basic Linear Algebra using native CUBLAS library+Description:+ Basic Linear Algebra using native CUBLAS library.+ Currently only support for the most basic batched matrix operations.+Tested-With: GHC==7.8.3+Cabal-Version: >=1.14+Build-Type: Simple++Flag buildExamples+ description: Build example executables+ default: False++Source-Repository this+ Tag: 0.0+ Type: darcs+ Location: http://code.haskell.org/~thielema/accelerate-cublas/++Source-Repository head+ Type: darcs+ Location: http://code.haskell.org/~thielema/accelerate-cublas/++Library+ Build-Depends:+ accelerate-arithmetic >=0.0.1 && <0.1,+ accelerate-utility >=0.1 && <0.2,+ accelerate-cuda >=0.15 && <0.16,+ accelerate-io >=0.15 && <0.16,+ accelerate >=0.15 && <0.16,+ cublas >=0.2.0.2 && <0.3,+ cuda >=0.5 && <0.7,+ vector >=0.10.11 && <0.11,+ utility-ht >=0.0.8 && <0.1,+ base >=4.5 && <4.8++ GHC-Options: -Wall -fwarn-missing-import-lists+ Hs-Source-Dirs: src+ Default-Language: Haskell98+ Exposed-Modules:+ Data.Array.Accelerate.CUBLAS.Level2.Batched+ Data.Array.Accelerate.CUBLAS.Level3.Batched+ Other-Modules:+ Data.Array.Accelerate.CUBLAS.Level3.Batched.Foreign++Executable accelerate-cublas-demo+ GHC-Options: -Wall -fwarn-missing-import-lists+ Hs-Source-Dirs: example+ Default-Language: Haskell98+ Main-Is: Main.hs+ Build-Depends:+ accelerate-cublas,+ accelerate-cuda,+ accelerate-arithmetic,+ accelerate-utility,+ accelerate,+ cublas,+ random >=1.0 && <1.1,+ base >=4.5 && <4.8
+ example/Main.hs view
@@ -0,0 +1,109 @@+module Main where++import qualified Data.Array.Accelerate.CUBLAS.Level3.Batched as Batched++import qualified Data.Array.Accelerate.Arithmetic.LinearAlgebra as ALinAlg++import qualified Data.Array.Accelerate.Utility.Lift.Exp as Exp+import Data.Array.Accelerate.Utility.Lift.Exp (expr)++import Data.Array.Accelerate (DIM1, Z(Z), (:.)((:.)), (!), (?), (==*))+import qualified Data.Array.Accelerate.CUDA as AC+import qualified Data.Array.Accelerate as A++import qualified Foreign.CUDA.Cublas as Cublas++import System.Random (randomRs, mkStdGen)+++factorMatrices :: (ALinAlg.Matrix DIM1 Double, ALinAlg.Matrix DIM1 Double)+factorMatrices =+ let numMats = 100+ f =+ Exp.modify (expr :. expr :. expr :. expr) $+ \(_z :. i :. j :. k) -> A.fromIntegral (i+j+k)+ in (A.generate (A.constant $ Z :. numMats :. 3 :. 4) f,+ A.generate (A.constant $ Z :. numMats :. 4 :. 2) f)++mainMul :: Cublas.Handle -> IO ()+mainMul handle = do+ print factorMatrices+ print $ AC.run $+ case factorMatrices of+ (a,b) -> Batched.mul handle 1 a b+++luMatrices :: (ALinAlg.Matrix Z Double, ALinAlg.Matrix Z Double)+luMatrices =+ (A.use $ A.fromList (Z :. 4 :. 4) $+ 2 : 0 : 0 : 0 :+ 1 : 3 : 0 : 0 :+ 0 : 1 : 4 : 0 :+ 0 : 0 : 1 : 5 :+ [],+ A.use $ A.fromList (Z :. 4 :. 4) $+ 0 : 1 : 1 : 0 :+ 0 : 0 : 1 : 1 :+ 0 : 0 : 0 : 1 :+ 1 : 1 : 0 : 0 :+ [])++permMatrix :: ALinAlg.Matrix Z Double+permMatrix =+ A.use $ A.fromList (Z :. 4 :. 4) $+ 0 : 2 : 0 : 0 :+ 0 : 0 : 0 : 4 :+ 0 : 0 : 3 : 0 :+ 1 : 0 : 0 : 0 :+ []++rhsMatrix :: ALinAlg.Matrix Z Double+rhsMatrix =+ A.use $ A.fromList (Z :. 4 :. 2) $+ 2 : 5 :+ 8 : 6 :+ 8 : 3 :+ 7 : 6 :+ []++append ::+ (A.Elt a) =>+ ALinAlg.Matrix Z a -> ALinAlg.Matrix Z a -> ALinAlg.Matrix DIM1 a+append x y =+ let (_z:.m:.n) =+ Exp.unlift (expr:.expr:.expr) $ A.intersect (A.shape x) (A.shape y)+ in A.generate (A.lift $ Z :. (2::Int) :. m :. n) $+ Exp.modify (expr :. expr :. expr :. expr) $+ \(_z :. k :. i :. j) ->+ let ix = A.index2 i j+ in k==*0 ? (x!ix, y!ix)++mainLU :: Cublas.Handle -> IO ()+mainLU handle = do+ print luMatrices+ let mat =+ append permMatrix $+ case luMatrices of+ (a,b) -> Batched.mul handle 1 a b+ lu = Batched.lu handle mat+ print $ AC.run $ Batched.luSolve handle lu $+ Batched.mul handle 1 mat $+ A.replicate (A.lift $ Z :. (2::Int) :. A.All :. A.All) rhsMatrix++mainInv :: Cublas.Handle -> IO ()+mainInv handle = do+ let dim = 4+ mat =+ A.fromList (Z:.3:.dim:.dim :: A.DIM3) $+ randomRs (-1,1::Float) $ mkStdGen 42+ test a =+ let (inv, info) = Batched.inv handle a+ in A.lift (Batched.mul handle 1 a inv, info)+ print $ AC.run1 test mat++main :: IO ()+main = do+ handle <- Cublas.create+ mainMul handle+ mainLU handle+ mainInv handle
+ src/Data/Array/Accelerate/CUBLAS/Level2/Batched.hs view
@@ -0,0 +1,42 @@+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+module Data.Array.Accelerate.CUBLAS.Level2.Batched (+ Level3.Element,+ mul,+ mac,+ ) where++import qualified Data.Array.Accelerate.CUBLAS.Level3.Batched as Level3+import Data.Array.Accelerate.CUBLAS.Level3.Batched (Element)++import qualified Data.Array.Accelerate.Arithmetic.LinearAlgebra as ALinAlg++import qualified Data.Array.Accelerate as A+import Data.Array.Accelerate (Exp)++import qualified Foreign.CUDA.Cublas as Cublas+++mul ::+ (A.Shape ix, A.Slice ix, Eq ix, Element a, A.Elt a, A.IsNum a) =>+ Cublas.Handle ->+ Exp a ->+ ALinAlg.Matrix ix a -> ALinAlg.Vector ix a ->+ ALinAlg.Vector ix a+mul handle alpha a b =+ ALinAlg.vectorFromColumn $+ Level3.mul handle alpha a (ALinAlg.columnFromVector b)++mac ::+ (A.Shape ix, A.Slice ix, Eq ix, Element a, A.Elt a, A.IsNum a) =>+ Cublas.Handle ->+ Exp a -> ALinAlg.Matrix ix a -> ALinAlg.Vector ix a ->+ Exp a -> ALinAlg.Vector ix a ->+ ALinAlg.Vector ix a+mac handle alpha a b beta c =+ A.reshape (A.shape c) $+ Level3.mac handle+ alpha a (ALinAlg.columnFromVector b)+ beta (ALinAlg.columnFromVector c)
+ src/Data/Array/Accelerate/CUBLAS/Level3/Batched.hs view
@@ -0,0 +1,266 @@+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+module Data.Array.Accelerate.CUBLAS.Level3.Batched (+ Element,+ mul,+ mac,+ LU,+ lu,+ luInv,+ inv,+ luSolve,+ newtonInverseStep,+ newtonInverse,+ ) where++import qualified Data.Array.Accelerate.CUBLAS.Level3.Batched.Foreign as Foreign+import Data.Array.Accelerate.CUBLAS.Level3.Batched.Foreign (Element, Vector)++import qualified Data.Array.Accelerate.Arithmetic.LinearAlgebra as ALinAlg++import qualified Data.Array.Accelerate.Utility.Sliced1 as Sliced1+import qualified Data.Array.Accelerate.Utility.Sliced as Sliced+import qualified Data.Array.Accelerate.Utility.Arrange as Arrange+import qualified Data.Array.Accelerate.Utility.Loop as Loop+import qualified Data.Array.Accelerate.Utility.Lift.Acc as Acc+import qualified Data.Array.Accelerate.Utility.Lift.Exp as Exp+import Data.Array.Accelerate.Utility.Lift.Acc (acc, expr)++import Data.Array.Accelerate (Exp, (:.)((:.)), (!))+import qualified Data.Array.Accelerate.CUDA.Foreign as AF+import qualified Data.Array.Accelerate.IO as AIO+import qualified Data.Array.Accelerate as A++import qualified Foreign.CUDA.Cublas as Cublas++import qualified Data.Vector.Storable as V+import qualified Data.Vector.Storable.Mutable as MV++import Control.Monad.ST (ST)+import Control.Monad (zipWithM_)++import Data.Tuple.HT (mapSnd, uncurry3)++import Data.Word (Word32)+++mul ::+ (A.Shape ix, A.Slice ix, Eq ix, Element a, A.Elt a, A.IsNum a) =>+ Cublas.Handle ->+ Exp a ->+ ALinAlg.Matrix ix a -> ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a+mul handle alpha a b =+ A.foreignAcc+ (AF.CUDAForeignAcc "mul" $ uncurry3 $ Foreign.mul handle)+ (Acc.modify (expr,acc,acc) $ \(alpha0, a0, b0) ->+ A.map (alpha0 *) $+ ALinAlg.multiplyMatrixMatrix a0 b0)+ $+ A.lift (A.unit alpha, a, b)++mac ::+ (A.Shape ix, A.Slice ix, Eq ix, Element a, A.Elt a, A.IsNum a) =>+ Cublas.Handle ->+ Exp a -> ALinAlg.Matrix ix a -> ALinAlg.Matrix ix a ->+ Exp a -> ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a+mac handle alpha a b beta c =+ A.foreignAcc+ (AF.CUDAForeignAcc "mac" $+ \((alpha0, a0, b0), (beta0, c0)) ->+ Foreign.mac handle alpha0 a0 b0 beta0 c0)+ (Acc.modify ((expr,acc,acc),(expr,acc)) $+ \((alpha0, a0, b0), (beta0, c0)) ->+ A.zipWith (+)+ (A.map (alpha0 *) $+ ALinAlg.multiplyMatrixMatrix a0 b0)+ (A.map (beta0 *) c0))+ $+ A.lift ((A.unit alpha, a, b), (A.unit beta, c))++++newtype LU ix a =+ LU {+ _getLU ::+ (ALinAlg.Matrix ix a,+ ALinAlg.Vector ix Word32, ALinAlg.Scalar ix Word32)+ }++lu ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ ALinAlg.Matrix ix a -> LU ix a+lu handle =+ LU . A.unlift . cudaAcc "lu" (Foreign.lu handle)+++luInv ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ LU ix a -> (ALinAlg.Matrix ix a, ALinAlg.Scalar ix Word32)+luInv handle (LU sol@(_,_,info)) =+ (cudaAcc "luInv" (Foreign.luInv handle) $ A.lift sol, info)+++{- |+Returns the inverted matrix and a rank information.+If the matrix is invertible, then the rank information is zero.+Otherwise it is the matrix rank plus 1.+-}+inv, _inv ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ ALinAlg.Matrix ix a ->+ (ALinAlg.Matrix ix a, ALinAlg.Scalar ix Word32)+inv handle = luInv handle . lu handle++{- |+maximum size of matrices is 32 in CUDA-6.0 and CUDA-6.5.+-}+_inv handle =+ A.unlift . cudaAcc "inv" (Foreign.inv handle)+++{- |+Matrices with sizes larger than 32+are only supported starting with CUDA-6.5.+In CUDA-6.0 you will get the error+@CUBLAS Exception: unsupported value or parameter passed to a function@.+On CUDA-6.0 you may prefer 'luInv' which works surprisingly.+-}+luSolve ::+ (A.Shape ix, A.Slice ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ LU ix a ->+ ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a+luSolve handle (LU (luMat, pivots, _info)) =+ let perm = permutationFromPivotsAcc $ A.map (subtract 1) pivots+ in applyRowPerm perm+ .+ cudaAcc "luSolve" (uncurry $ Foreign.luSolve handle $ Acc.singleton 1)+ .+ A.lift . (,) luMat+++_applyColPerm ::+ (A.Shape ix, A.Slice ix, A.Elt a) =>+ ALinAlg.Vector ix Word32 ->+ ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a+_applyColPerm perm arr =+ Arrange.mapWithIndex+ (Exp.modify2 (expr:.expr:.expr) expr $+ \(ix :. j :. _i) src -> arr ! A.lift (ix :. j :. src)) $+ A.replicate (A.lift $ A.Any :. Sliced1.length arr :. A.All) $+ A.map (A.fromIntegral :: Exp Word32 -> Exp Int) perm++applyRowPerm ::+ (A.Shape ix, A.Slice ix, A.Elt a) =>+ ALinAlg.Vector ix Word32 ->+ ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a+applyRowPerm perm arr =+ Arrange.mapWithIndex+ (Exp.modify2 (expr:.expr:.expr) expr $+ \(ix :. _j :. i) src -> arr ! A.lift (ix :. src :. i)) $+ A.replicate (A.lift $ A.Any :. Sliced.length arr) $+ A.map (A.fromIntegral :: Exp Word32 -> Exp Int) perm++permutationFromPivotsAcc ::+ (A.Shape ix) =>+ ALinAlg.Vector ix Word32 -> ALinAlg.Vector ix Word32+permutationFromPivotsAcc =+ cudaAcc "permutations" $ \arr -> do+ AF.peekArray arr+ let perm = permutationFromPivots arr+ AF.useArray perm+ return perm++permutationFromPivots ::+ (A.Shape ix) =>+ Vector ix Word32 -> Vector ix Word32+permutationFromPivots vec =+ let sh = A.arrayShape vec+ in AIO.fromVectors sh $+ mapSnd (permutationsFromPivotsSlices sh) $+ AIO.toVectors vec++permutationsFromPivotsSlices ::+ (A.Shape sh) =>+ sh :. Int -> V.Vector Word32 -> V.Vector Word32+permutationsFromPivotsSlices (shape:.width) pivots = V.create (do+ perm <- MV.new $ V.length pivots+ mapM_+ (\k ->+ permutationFromPivotsMutableBackward+ (V.slice k width pivots)+ (MV.slice k width perm))+ (take (A.arraySize shape) [0, width ..])+ return perm)++{- |+works always, but requires two traversals through the array+-}+_permutationFromPivotsMutable ::+ V.Vector Word32 -> MV.MVector s Word32 -> ST s ()+_permutationFromPivotsMutable pivots perm = do+ let ixs = V.enumFromN 0 (V.length pivots)+ V.copy perm ixs+ V.zipWithM_+ (\k j -> MV.swap perm (fromIntegral k) (fromIntegral j))+ (V.reverse ixs) (V.reverse pivots)++-- | works only if forall i. pivot!!i >= i+permutationFromPivotsMutableBackward ::+ V.Vector Word32 -> MV.MVector s Word32 -> ST s ()+permutationFromPivotsMutableBackward pivots perm = do+ zipWithM_+ (\k j -> do+ MV.write perm k (fromIntegral k)+ MV.swap perm k (fromIntegral j))+ (iterate (subtract 1) $ V.length pivots - 1) (V.toList $ V.reverse pivots)++-- | works only if forall i. pivot!!i <= i+_permutationFromPivotsMutableForward ::+ V.Vector Word32 -> MV.MVector s Word32 -> ST s ()+_permutationFromPivotsMutableForward pivots perm = do+ zipWithM_+ (\k j -> do+ MV.write perm k (fromIntegral k)+ MV.swap perm k (fromIntegral j))+ [0..] (V.toList pivots)++++newtonInverseStep ::+ (A.Shape ix, A.Slice ix, Eq ix, Element a, A.Elt a, A.IsNum a) =>+ Cublas.Handle ->+ ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a+newtonInverseStep h a x =+ mac h (-1) x (mul h 1 a x) 2 x++newtonInverse ::+ (A.Shape ix, A.Slice ix, Eq ix, Element a, A.Elt a, A.IsNum a) =>+ Cublas.Handle ->+ A.Exp Int ->+ ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a ->+ ALinAlg.Matrix ix a+newtonInverse h n seed a =+ Loop.nest n (newtonInverseStep h a) seed+++cudaAcc ::+ (A.Arrays res, A.Arrays acc) =>+ String -> (acc -> AF.CIO res) -> A.Acc acc -> A.Acc res+cudaAcc name f =+ A.foreignAcc+ (AF.CUDAForeignAcc name f)+ (error $ name ++ ": requires CUDA backend")
+ src/Data/Array/Accelerate/CUBLAS/Level3/Batched/Foreign.hs view
@@ -0,0 +1,233 @@+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+module Data.Array.Accelerate.CUBLAS.Level3.Batched.Foreign where++import Data.Array.Accelerate.Array.Sugar (EltRepr)+import Data.Array.Accelerate (Array, Shape, Z(Z), (:.)((:.)))+import qualified Data.Array.Accelerate.CUDA.Foreign as AF+import qualified Data.Array.Accelerate as A++import qualified Foreign.CUDA.Cublas as Cublas+import Foreign.CUDA.Ptr (DevicePtr, castDevPtr, advanceDevPtr)++import Foreign.C.Types (CFloat, CDouble)+import Foreign.Storable (Storable)++import Data.Word (Word32)+++type Matrix ix = Array (ix :. Int :. Int)+type Vector ix = Array (ix :. Int)+type Scalar ix = Array ix+++mul ::+ (Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ A.Scalar a -> Matrix ix a -> Matrix ix a ->+ AF.CIO (Matrix ix a)+mul handle alpha a b = do+ let (aNumMatrices :. n :. ak) = A.arrayShape a+ let (bNumMatrices :. bk :. m) = A.arrayShape b+ let k = unify "mul: matrix sizes mismatch" ak bk+ let numMatrices =+ unify "mul: mismatching shapes of matrix arrays"+ aNumMatrices bNumMatrices+ c <- AF.allocateArray (numMatrices :. n :. m)+ (pas, lda) <- arrayPtrs a+ (pbs, ldb) <- arrayPtrs b+ (pcs, ldc) <- arrayPtrs c+ AF.liftIO $+ Cublas.gemmBatched handle Cublas.N Cublas.N m n k+ (storableFromScalar alpha)+ pbs ldb+ pas lda+ 0+ pcs ldc+ (A.arraySize numMatrices)+ return c++mac ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ A.Scalar a -> Matrix ix a -> Matrix ix a ->+ A.Scalar a -> Matrix ix a ->+ AF.CIO (Matrix ix a)+mac handle alpha a b beta c = do+ let (aNumMatrices :. an :. bk) = A.arrayShape a+ let (bNumMatrices :. ak :. bm) = A.arrayShape b+ let (cNumMatrices :. cn :. cm) = A.arrayShape c+ let k = unify "mac: matrix sizes mismatch" ak bk+ let n = unify "mac: matrix sizes mismatch" an cn+ let m = unify "mac: matrix sizes mismatch" bm cm+ let numMatrices =+ let msg = "mac: mismatching shapes of matrix arrays"+ in unify msg aNumMatrices $+ unify msg bNumMatrices cNumMatrices+ d <- AF.allocateArray (numMatrices :. n :. m)+ AF.copyArray c d+ (pas, lda) <- arrayPtrs a+ (pbs, ldb) <- arrayPtrs b+ (pds, ldd) <- arrayPtrs d+ AF.liftIO $+ Cublas.gemmBatched handle Cublas.N Cublas.N m n k+ (storableFromScalar alpha)+ pbs ldb+ pas lda+ (storableFromScalar beta)+ pds ldd+ (A.arraySize numMatrices)+ return d++lu ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ Matrix ix a ->+ AF.CIO (Matrix ix a, Vector ix Word32, Scalar ix Word32)+lu handle a = do+ let sh@(numMatrices :. n :. k) = A.arrayShape a+ let size = unify "lu: matrices must have square shape" n k+ b <- AF.allocateArray sh+ AF.copyArray a b+ (pbs, ldb) <- arrayPtrs b++ pivot <- AF.allocateArray (numMatrices :. size)+ pivotPtr <- devicePtrsOfArray pivot++ info <- AF.allocateArray numMatrices+ infoPtr <- devicePtrsOfArray info++ AF.liftIO $+ Cublas.getrfBatched handle size+ pbs ldb+ pivotPtr infoPtr+ (A.arraySize numMatrices)+ return (b, pivot, info)++luInv ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ (Matrix ix a, Vector ix Word32, Scalar ix Word32) ->+ AF.CIO (Matrix ix a)+luInv handle (a, pivot, info) = do+ let sh@(numMatrices :. n :. k) = A.arrayShape a+ let size = unify "luInv: matrices must have square shape" n k+ c <- AF.allocateArray sh+ AF.copyArray a c+ (pas, lda) <- arrayPtrs a+ (pcs, ldc) <- arrayPtrs c++ pivotPtr <- devicePtrsOfArray pivot+ infoPtr <- devicePtrsOfArray info++ AF.liftIO $+ Cublas.getriBatched handle size+ pas lda+ pivotPtr+ pcs ldc+ infoPtr+ (A.arraySize numMatrices)+ return c++luSolve ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ A.Scalar a ->+ Matrix ix a ->+ Matrix ix a ->+ AF.CIO (Matrix ix a)+luSolve handle alpha a b = do+ let (aNumMatrices :. an :. ak) = A.arrayShape a+ let sh@(bNumMatrices :. bk :. m) = A.arrayShape b+ let n =+ unify "luSolve: matrices must have square shape" an $+ unify "luSolve: matrices dimensions must match" ak bk+ let count =+ A.arraySize $+ unify "mul: mismatching shapes of matrix arrays"+ aNumMatrices bNumMatrices+ c <- AF.allocateArray sh+ AF.copyArray b c+ (pas, lda) <- arrayPtrs a+ (pcs, ldc) <- arrayPtrs c++ AF.liftIO $ do+ Cublas.trsmBatched handle+ Cublas.SideRight Cublas.Upper Cublas.N Cublas.NonUnit m n+ (storableFromScalar alpha)+ pas lda+ pcs ldc+ count+ Cublas.trsmBatched handle+ Cublas.SideRight Cublas.Lower Cublas.N Cublas.Unit m n+ (storableFromScalar alpha)+ pas lda+ pcs ldc+ count+ return c++inv ::+ (A.Shape ix, Eq ix, Element a, A.Elt a) =>+ Cublas.Handle ->+ Matrix ix a ->+ AF.CIO (Matrix ix a, Scalar ix Word32)+inv handle a = do+ let sh@(numMatrices :. n :. k) = A.arrayShape a+ let size = unify "inv: matrices must have square shape" n k+ b <- AF.allocateArray sh+ (pas, lda) <- arrayPtrs a+ (pbs, ldb) <- arrayPtrs b+ info <- AF.allocateArray numMatrices+ infoPtr <- fmap (castDevPtr . snd) $ AF.devicePtrsOfArray info++ AF.liftIO $+ Cublas.matinvBatched handle size+ pas lda+ pbs ldb+ infoPtr+ (A.arraySize numMatrices)+ return (b, info)+++type Element a =+ (AF.DevicePtrs (EltRepr a) ~ ((), DevicePtr a),+ Fractional (StorableOf a),+ Cublas.Cublas (StorableOf a),+ Storable (StorableOf a),+ Real a)++type family StorableOf float+type instance StorableOf Float = CFloat+type instance StorableOf Double = CDouble++storableFromScalar ::+ (Real a, StorableOf a ~ b, Fractional b) => A.Scalar a -> b+storableFromScalar x = realToFrac $ A.indexArray x Z++genPointers ::+ (Storable a) =>+ Int -> Int -> DevicePtr a -> [DevicePtr a]+genPointers n size p =+ take n $ iterate (flip advanceDevPtr size) p++arrayPtrs ::+ (A.Shape ix,+ Storable a, StorableOf e ~ a,+ AF.DevicePtrs (EltRepr e) ~ ((), DevicePtr e)) =>+ Matrix ix e -> AF.CIO ([DevicePtr a], Int)+arrayPtrs arr = do+ let (numMatrices :. n :. k) = A.arrayShape arr+ pa <- devicePtrsOfArray arr+ return (genPointers (A.arraySize numMatrices) (n*k) pa, k)++devicePtrsOfArray ::+ (A.Shape ix, AF.DevicePtrs (EltRepr e) ~ ((), DevicePtr e)) =>+ Scalar ix e -> AF.CIO (DevicePtr a)+devicePtrsOfArray arr = do+ ((), pa) <- AF.devicePtrsOfArray arr+ return $ castDevPtr pa++unify :: (Eq a) => String -> a -> a -> a+unify msg a b = if a == b then a else error msg