add: filter after multiplication

Author: IgorErin

src/GraphBLAS-sharp.Backend/Common/ClArray.fs

@@ -389,7 +389,7 @@ module ClArray =
 
             result
 
-    let private assignOption2 (clContext: ClContext) workGroupSize (op: Expr<'a -> 'b -> 'c option>) =
+    let assignOption2 (clContext: ClContext) workGroupSize (op: Expr<'a -> 'b -> 'c option>) =
 
         let assign =
             <@ fun (ndRange: Range1D) length (firstValues: ClArray<'a>) (secondValues: ClArray<'b>) (positions: ClArray<int>) (result: ClArray<'c>) resultLength ->

src/GraphBLAS-sharp.Backend/Matrix/CSRMatrix/SpGEMM/Expand.fs

@@ -9,6 +9,7 @@ open GraphBLAS.FSharp.Backend.Objects.ClContext
 open GraphBLAS.FSharp.Backend.Objects.ArraysExtensions
 open GraphBLAS.FSharp.Backend.Objects
 open GraphBLAS.FSharp.Backend.Objects.ClCell
+open FSharp.Quotations
 
 type Indices = ClArray<int>
 
@@ -65,7 +66,40 @@ module Expand =
 
             length, segmentsLengths
 
-    let expand (clContext: ClContext) workGroupSize opMul =
+    let multiply (clContext: ClContext) workGroupSize (predicate: Expr<'a -> 'b -> 'c option>) =
+        let getBitmap =
+            ClArray.map2<'a, 'b, int> clContext workGroupSize
+            <| Map.chooseBitmap2 predicate
+
+        let prefixSum = PrefixSum.standardExcludeInplace clContext workGroupSize
+
+        let assignValues = ClArray.assignOption2 clContext workGroupSize predicate
+
+        let scatter = Scatter.lastOccurrence clContext workGroupSize // TODO(last ?)
+
+        fun (processor: MailboxProcessor<_>) (firstValues: ClArray<'a>) (secondValues: ClArray<'b>) (columns: Indices) (rows: Indices) ->
+
+            let positions = getBitmap processor DeviceOnly firstValues secondValues
+
+            let resultLength =
+                (prefixSum processor positions)
+                    .ToHostAndFree(processor)
+
+            let resultColumns = clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
+
+            scatter processor positions columns resultColumns
+
+            let resultRows = clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
+
+            scatter processor positions rows resultRows
+
+            let resultValues = clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
+
+            assignValues processor firstValues secondValues positions resultValues
+
+            resultValues, resultColumns, resultRows
+
+    let expand (clContext: ClContext) workGroupSize =
 
         let idScatter = Scatter.initLastOccurrence Map.id clContext workGroupSize
 
@@ -89,8 +123,6 @@ module Expand =
 
         let BGather = Gather.run clContext workGroupSize
 
-        let mul = ClArray.map2 clContext workGroupSize opMul
-
         fun (processor: MailboxProcessor<_>) lengths (segmentsPointers: Indices) (leftMatrix: ClMatrix.CSR<'a>) (rightMatrix: ClMatrix.CSR<'b>) ->
 
             // Compute A positions
@@ -150,13 +182,8 @@ module Expand =
 
             BPositions.Free processor
 
-            // multiply values TODO(filter values)
-            let values = mul processor DeviceOnly AValues BValues
-
-            AValues.Free processor
-            BValues.Free processor
-
-            values, columns, rows
+            // left, right matrix values, columns and rows indices
+            AValues, BValues, columns, rows
 
     let sortByColumnsAndRows (clContext: ClContext) workGroupSize =
 
@@ -227,7 +254,9 @@ module Expand =
 
         let getSegmentPointers = getSegmentPointers clContext workGroupSize
 
-        let expand = expand clContext workGroupSize opMul
+        let expand = expand clContext workGroupSize
+
+        let multiply = multiply clContext workGroupSize opMul
 
         let sort = sortByColumnsAndRows clContext workGroupSize
 
@@ -237,24 +266,37 @@ module Expand =
 
             let length, segmentPointers = getSegmentPointers processor leftMatrix rightMatrix
 
-            let values, columns, rows =
+            // expand
+            let leftMatrixValues, rightMatrixValues, columns, rows =
                 expand processor length segmentPointers leftMatrix rightMatrix
 
-            printfn $"expanded values: %A{values.ToHost processor}"
+            printfn $"left matrix values: %A{leftMatrixValues.ToHost processor}"
+            printfn $"right matrix values: %A{rightMatrixValues.ToHost processor}"
             printfn $"expanded columns: %A{columns.ToHost processor}"
             printfn $"expanded rows: %A{rows.ToHost processor}"
 
+            // multiply
+            let resultValues, resultColumns, resultRows =
+                multiply processor leftMatrixValues rightMatrixValues columns rows
+
+            leftMatrixValues.Free processor
+            rightMatrixValues.Free processor
+            columns.Free processor
+            rows.Free processor
+
+            // sort
             let sortedValues, sortedColumns, sortedRows =
-                sort processor values columns rows
+                sort processor resultValues resultColumns resultRows
 
             printfn $"sorted values: %A{sortedValues.ToHost processor}"
             printfn $"sorted columns: %A{sortedColumns.ToHost processor}"
             printfn $"sorted rows: %A{sortedRows.ToHost processor}"
 
-            values.Free processor
-            columns.Free processor
-            rows.Free processor
+            resultValues.Free processor
+            resultColumns.Free processor
+            resultRows.Free processor
 
+            // addition
             let reducedValues, reducedColumns, reducedRows =
                 reduce processor allocationMode sortedValues sortedColumns sortedRows
 

src/GraphBLAS-sharp.Backend/Quotes/Arithmetic.fs

@@ -141,10 +141,20 @@ module ArithmeticOperations =
 
     let inline createPair zero op opQ = binOpQ zero opQ, binOp zero op
 
+    // addition
     let intAdd = createPair 0 (+) <@ (+) @>
 
     let boolAdd = createPair false (||) <@ (||) @>
 
     let floatAdd = createPair 0.0 (+) <@ (+) @>
 
     let float32Add = createPair 0.0f (+) <@ (+) @>
+
+    // multiplication
+    let intMul = createPair 0 (*) <@ (*) @>
+
+    let boolMul = createPair false (&&) <@ (&&) @>
+
+    let floatMul = createPair 0.0 (*) <@ (*) @>
+
+    let float32Mul = createPair 0.0f (*) <@ (*) @>

tests/GraphBLAS-sharp.Tests/Common/ClArray/Choose.fs

@@ -85,7 +85,7 @@ let createTest2 (isEqual: 'a -> 'a -> bool) (opMapQ, opMap) testFun =
     let testFun = testFun context Utils.defaultWorkGroupSize opMapQ
 
     makeTest2 isEqual opMap testFun
-    |> testPropertyWithConfig { config with maxTest = 1000 } $"test on %A{typeof<'a>}"
+    |> testPropertyWithConfig config $"test on %A{typeof<'a>}"
 
 let tests2 =
     [ createTest2 (=) ArithmeticOperations.intAdd ClArray.choose2

tests/GraphBLAS-sharp.Tests/Helpers.fs

@@ -229,7 +229,7 @@ module HostPrimitives =
 
         result
 
-    let array2DMultiplication mul add leftArray rightArray =
+    let array2DMultiplication zero mul add leftArray rightArray =
         if Array2D.length2 leftArray <> Array2D.length1 rightArray then
             failwith "Incompatible matrices"
 
@@ -239,7 +239,8 @@ module HostPrimitives =
         <| fun i j ->
                 (leftArray.[i, *], rightArray.[*, j])
                 ||> Array.map2 mul
-                |> Array.reduce add
+                |> Array.choose id
+                |> Array.fold add zero
 
 module Context =
     type TestContext =

tests/GraphBLAS-sharp.Tests/Matrix/SpGeMM.fs

@@ -2,6 +2,7 @@ module GraphBLAS.FSharp.Tests.Matrix.SpGeMM
 
 open Expecto
 open GraphBLAS.FSharp.Backend.Matrix.CSR.SpGeMM
+open GraphBLAS.FSharp.Backend.Quotes
 open GraphBLAS.FSharp.Test
 open Microsoft.FSharp.Collections
 open GraphBLAS.FSharp.Backend
@@ -76,7 +77,7 @@ let getSegmentsTests =
       createTest ((=) 0uy) Expand.getSegmentPointers ]
     |> testList "get segment pointers"
 
-let expand length segmentPointers mulOp (leftMatrix: Matrix.CSR<'a>) (rightMatrix: Matrix.CSR<'b>) =
+let expand length segmentPointers (leftMatrix: Matrix.CSR<'a>) (rightMatrix: Matrix.CSR<'b>) =
     let extendPointers pointers =
         Array.pairwise pointers
         |> Array.map (fun (fst, snd) -> snd - fst)
@@ -106,11 +107,9 @@ let expand length segmentPointers mulOp (leftMatrix: Matrix.CSR<'a>) (rightMatri
         |> Array.concat
         |> Array.unzip
 
-    let expectedValues = Array.map2 mulOp leftMatrixValues rightMatrixValues
+    leftMatrixValues, rightMatrixValues, expectedColumns, expectedRows
 
-    expectedValues, expectedColumns, expectedRows
-
-let makeExpandTest isEqual zero opMul testFun (leftArray: 'a [,], rightArray: 'a [,]) =
+let makeExpandTest isEqual zero testFun (leftArray: 'a [,], rightArray: 'a [,]) =
 
         let leftMatrix = createCSRMatrix leftArray <| isEqual zero
 
@@ -126,51 +125,55 @@ let makeExpandTest isEqual zero opMul testFun (leftArray: 'a [,], rightArray: 'a
             let clRightMatrix = rightMatrix.ToDevice context
             let clSegmentPointers = context.CreateClArray segmentPointers
 
-            let (clActualValues: ClArray<'a>), (clActualColumns: ClArray<int>), (clActualRows: ClArray<int>) =
+            let (clActualLeftValues: ClArray<'a>), (clActualRightValues: ClArray<'a>), (clActualColumns: ClArray<int>), (clActualRows: ClArray<int>) =
                 testFun processor length clSegmentPointers clLeftMatrix clRightMatrix
 
             clLeftMatrix.Dispose processor
             clRightMatrix.Dispose processor
             clSegmentPointers.Free processor
 
-            let actualValues = clActualValues.ToHostAndFree processor
+            let actualLeftValues = clActualLeftValues.ToHostAndFree processor
+            let actualRightValues = clActualRightValues.ToHostAndFree processor
             let actualColumns = clActualColumns.ToHostAndFree processor
             let actualRows = clActualRows.ToHostAndFree processor
 
-            let expectedValues, expectedColumns, expectedRows =
-                expand length segmentPointers opMul leftMatrix rightMatrix
+            let expectedLeftMatrixValues, expectedRightMatrixValues, expectedColumns, expectedRows =
+                expand length segmentPointers leftMatrix rightMatrix
+
+            "Left values must be the same"
+            |> Utils.compareArrays isEqual actualLeftValues expectedLeftMatrixValues
 
-            "Values must be the same"
-            |> Utils.compareArrays isEqual actualValues expectedValues
+            "Right values must be the same"
+            |> Utils.compareArrays isEqual actualRightValues expectedRightMatrixValues
 
             "Columns must be the same"
             |> Utils.compareArrays (=) actualColumns expectedColumns
 
             "Rows must be the same"
             |> Utils.compareArrays (=) actualRows expectedRows
 
-let createExpandTest isEqual (zero: 'a) opMul opMulQ testFun =
+let createExpandTest isEqual (zero: 'a) testFun =
 
-    let testFun = testFun context Utils.defaultWorkGroupSize opMulQ
+    let testFun = testFun context Utils.defaultWorkGroupSize
 
-    makeExpandTest isEqual zero opMul testFun
+    makeExpandTest isEqual zero testFun
     |> testPropertyWithConfig config $"test on %A{typeof<'a>}"
 
 let expandTests =
-    [ createExpandTest (=) 0 (*) <@ (*) @> Expand.expand
+    [ createExpandTest (=) 0 Expand.expand
 
       if Utils.isFloat64Available context.ClDevice then
-        createExpandTest Utils.floatIsEqual 0.0 (*) <@ (*) @> Expand.expand
+        createExpandTest Utils.floatIsEqual 0.0 Expand.expand
 
-      createExpandTest Utils.float32IsEqual 0f (*) <@ (*) @> Expand.expand
-      createExpandTest (=) false (&&) <@ (&&) @> Expand.expand
-      createExpandTest (=) 0uy (*) <@ (*) @> Expand.expand ]
+      createExpandTest Utils.float32IsEqual 0f  Expand.expand
+      createExpandTest (=) false  Expand.expand
+      createExpandTest (=) 0uy  Expand.expand ]
     |> testList "Expand.expand"
 
 let checkGeneralResult zero isEqual actualValues actualColumns actualRows mul add (leftArray: 'a [,]) (rightArray: 'a [,]) =
 
     let expected =
-        HostPrimitives.array2DMultiplication mul add leftArray rightArray
+        HostPrimitives.array2DMultiplication zero mul add leftArray rightArray
         |> fun array -> Utils.createMatrixFromArray2D COO array (isEqual zero)
         |> function Matrix.COO matrix -> matrix | _ -> failwith "format miss"
 
@@ -217,15 +220,15 @@ let makeGeneralTest zero isEqual opMul opAdd testFun (leftArray: 'a [,], rightAr
                 checkGeneralResult zero isEqual actualValues actualColumns actualRows opMul opAdd leftArray rightArray
            with
            | ex when ex.Message = "InvalidBufferSize" -> ()
-           | ex -> raise ex
+           | _ -> reraise ()
 
-let createGeneralTest (zero: 'a) isEqual opAdd opAddQ opMul opMulQ testFun =
+let createGeneralTest (zero: 'a) isEqual opAddQ opAdd (opMulQ, opMul) testFun =
 
     let testFun = testFun context Utils.defaultWorkGroupSize opAddQ opMulQ
 
     makeGeneralTest zero isEqual opMul opAdd testFun
     |> testPropertyWithConfig { config with endSize = 10; maxTest = 1000 } $"test on %A{typeof<'a>}"
 
 let generalTests =
-    [ createGeneralTest 0 (=) (+) <@ (+) @> (*) <@ (*) @> Expand.run ]
+    [ createGeneralTest 0 (=) <@ (+) @> (+) ArithmeticOperations.intMul Expand.run ]
     |> testList "general"

tests/GraphBLAS-sharp.Tests/Program.fs

@@ -94,10 +94,11 @@ open GraphBLAS.FSharp.Tests.Matrix
 let allTests =
     testList
         "All tests"
-        [ // SpGeMM.generalTests
+        [ // SpGeMM.expandTests
+          SpGeMM.generalTests
           // Common.Gather.initTests
-          Common.ClArray.Choose.tests2 ]
-
+          // Common.ClArray.Choose.tests2 ]
+        ]
     |> testSequenced
 
 [<EntryPoint>]