From 2c7bf2d2417fdec96246580c3e45b2b2a9445967 Mon Sep 17 00:00:00 2001 From: dmjio Date: Sun, 14 Jun 2026 19:56:59 -0500 Subject: [PATCH] sparse: add sparseMatmul and tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit af_matmul already dispatches to an optimised sparse-dense kernel when the lhs is a CSR array; expose this as sparseMatmul in ArrayFire.Sparse so callers have a named, documented entry point without having to know that detail. Adds four hspec tests covering sparse×vector, sparse×matrix, identity, and agreement with the dense matmul path. Co-Authored-By: Claude Sonnet 4.6 --- src/ArrayFire/Sparse.hs | 29 +++++++++++++++++++++++++++++ test/ArrayFire/SparseSpec.hs | 29 +++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/src/ArrayFire/Sparse.hs b/src/ArrayFire/Sparse.hs index 888cd96..076e50b 100644 --- a/src/ArrayFire/Sparse.hs +++ b/src/ArrayFire/Sparse.hs @@ -26,6 +26,7 @@ import ArrayFire.Exception import ArrayFire.Types import ArrayFire.FFI import ArrayFire.Internal.Algorithm (af_any_true_all) +import ArrayFire.Internal.BLAS (af_matmul) import ArrayFire.Internal.Sparse import ArrayFire.Internal.Types import Data.Int @@ -352,3 +353,31 @@ sparseGetNNZ -> Int sparseGetNNZ a = fromIntegral (a `infoFromArray` af_sparse_get_nnz) + +-- | Multiply a sparse matrix by a dense matrix or vector. +-- +-- [ArrayFire Docs](http://arrayfire.org/docs/group__blas__func__matmul.htm) +-- +-- Calls the standard @af_matmul@ with a sparse left-hand side, which +-- ArrayFire dispatches to an optimised sparse-dense kernel. The sparse +-- input /must/ be in 'CSR' format; use 'sparseConvertTo' to convert first +-- if necessary. +-- +-- >>> let sp = createSparseArrayFromDense (matrix @Double (3,3) [[1,0,0],[0,2,0],[0,0,3]]) CSR +-- >>> sparseMatmul sp (vector @Double 3 [1,2,3]) +-- ArrayFire Array +-- [3 1 1 1] +-- 1.0000 +-- 4.0000 +-- 9.0000 +-- +sparseMatmul + :: AFType a + => Array a + -- ^ sparse CSR matrix (lhs) + -> Array a + -- ^ dense matrix or vector (rhs) + -> Array a + -- ^ dense result +sparseMatmul lhs rhs = + op2 lhs rhs (\p a b -> af_matmul p a b (toMatProp None) (toMatProp None)) diff --git a/test/ArrayFire/SparseSpec.hs b/test/ArrayFire/SparseSpec.hs index ec83520..470ff1f 100644 --- a/test/ArrayFire/SparseSpec.hs +++ b/test/ArrayFire/SparseSpec.hs @@ -5,6 +5,7 @@ import qualified ArrayFire as A import Control.Exception (evaluate) import Data.Int import Test.Hspec +import Test.Hspec.ApproxExpect (closeList) -- 3×3 diagonal matrix diag(1,2,3), stored column-major: -- col0=[1,0,0], col1=[0,2,0], col2=[0,0,3] @@ -48,3 +49,31 @@ spec = A.sparseGetNNZ sp `shouldBe` 3 A.sparseGetStorage sp `shouldBe` A.COO A.sparseToDense (A.sparseConvertTo sp A.CSR) `shouldBe` diag3 + + describe "sparseMatmul" $ do + let spDiag3 = A.createSparseArrayFromDense diag3 A.CSR + + it "sparse diagonal × dense vector: scales each element" $ do + -- diag(1,2,3) * [1,2,3]^T = [1,4,9]^T + let v = A.vector @Double 3 [1,2,3] + result = A.sparseMatmul spDiag3 v + closeList (A.toList result) [1, 4, 9] + + it "sparse diagonal × dense matrix: scales each row" $ do + -- diag(1,2,3) * [[1,4],[2,5],[3,6]] = [[1,4],[4,10],[9,18]] + -- column-major storage: col0=[1,2,3], col1=[4,5,6] + let b = A.mkArray @Double [3,2] [1,2,3,4,5,6] + result = A.sparseMatmul spDiag3 b + closeList (A.toList result) [1,4,9, 4,10,18] + + it "sparse identity × dense vector equals the vector" $ do + let eye3 = A.createSparseArrayFromDense + (A.mkArray @Double [3,3] [1,0,0, 0,1,0, 0,0,1]) A.CSR + v = A.vector @Double 3 [3,1,4] + closeList (A.toList (A.sparseMatmul eye3 v)) (A.toList v) + + it "sparseMatmul agrees with dense matmul" $ do + -- Any result from sparse × dense should equal the dense × dense product. + let dense = A.toList (A.matmul diag3 (A.mkArray @Double [3,2] [1,2,3,4,5,6]) A.None A.None) + sparse = A.toList (A.sparseMatmul spDiag3 (A.mkArray @Double [3,2] [1,2,3,4,5,6])) + closeList sparse dense