add: reduce by key option

Author: IgorErin

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

@@ -6,6 +6,7 @@ open Microsoft.FSharp.Control
 open Microsoft.FSharp.Quotations
 open GraphBLAS.FSharp.Backend.Objects.ClContext
 open GraphBLAS.FSharp.Backend.Objects.ArraysExtensions
+open GraphBLAS.FSharp.Backend.Objects.ClCell
 
 module Reduce =
     /// <summary>
@@ -616,3 +617,127 @@ module Reduce =
                 processor.Post(Msg.CreateRunMsg<_, _>(kernel))
 
                 firstReducedKeys, secondReducedKeys, reducedValues
+
+        /// <summary>
+        /// Reduces values by key. Each segment is reduced by one work item.
+        /// </summary>
+        /// <param name="clContext">ClContext.</param>
+        /// <param name="workGroupSize">Work group size.</param>
+        /// <param name="reduceOp">Operation for reducing values.</param>
+        /// <remarks>
+        /// The length of the result must be calculated in advance.
+        /// </remarks>
+        let segmentSequentialOption<'a> (clContext: ClContext) workGroupSize (reduceOp: Expr<'a -> 'a -> 'a option>) =
+
+            let kernel =
+                <@ fun (ndRange: Range1D) uniqueKeyCount keysLength (offsets: ClArray<int>) (firstKeys: ClArray<int>) (secondKeys: ClArray<int>) (values: ClArray<'a>) (reducedValues: ClArray<'a>) (firstReducedKeys: ClArray<int>) (secondReducedKeys: ClArray<int>) (resultPositions: ClArray<int>) ->
+
+                    let gid = ndRange.GlobalID0
+
+                    if gid < uniqueKeyCount then
+                        let startPosition = offsets.[gid]
+
+                        let firstSourceKey = firstKeys.[startPosition]
+                        let secondSourceKey = secondKeys.[startPosition]
+
+                        let mutable sum = Some values.[startPosition]
+
+                        let mutable currentPosition = startPosition + 1
+
+                        while currentPosition < keysLength
+                              && firstSourceKey = firstKeys.[currentPosition]
+                              && secondSourceKey = secondKeys.[currentPosition] do
+
+                            match sum with
+                            | Some value ->
+                                let result = ((%reduceOp) value values.[currentPosition]) // brahma error
+
+                                sum <- result
+                            | None ->
+                                sum <- Some values.[currentPosition]
+
+                            currentPosition <- currentPosition + 1
+
+                        match sum with
+                        | Some value ->
+                            reducedValues.[gid] <- value
+                            resultPositions.[gid] <- 1
+                        | None ->
+                            resultPositions.[gid] <- 0
+
+                        firstReducedKeys.[gid] <- firstSourceKey
+                        secondReducedKeys.[gid] <- secondSourceKey @>
+
+            let kernel = clContext.Compile kernel
+
+            let scatterData = Scatter.lastOccurrence clContext workGroupSize
+
+            let scatterIndices = Scatter.lastOccurrence clContext workGroupSize
+
+            let prefixSum = PrefixSum.standardExcludeInplace clContext workGroupSize
+
+            fun (processor: MailboxProcessor<_>) allocationMode (resultLength: int) (offsets: ClArray<int>) (firstKeys: ClArray<int>) (secondKeys: ClArray<int>) (values: ClArray<'a>) ->
+
+                let reducedValues =
+                    clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let firstReducedKeys =
+                    clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let secondReducedKeys =
+                    clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let resultPositions =
+                    clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
+
+                let ndRange =
+                    Range1D.CreateValid(resultLength, workGroupSize)
+
+                let kernel = kernel.GetKernel()
+
+                processor.Post(
+                    Msg.MsgSetArguments
+                        (fun () ->
+                            kernel.KernelFunc
+                                ndRange
+                                resultLength
+                                firstKeys.Length
+                                offsets
+                                firstKeys
+                                secondKeys
+                                values
+                                reducedValues
+                                firstReducedKeys
+                                secondReducedKeys
+                                resultPositions)
+                )
+
+                processor.Post(Msg.CreateRunMsg<_, _>(kernel))
+
+                let resultLength =
+                    (prefixSum processor resultPositions).ToHostAndFree processor
+
+                let resultValues =
+                    clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                scatterData processor resultPositions reducedValues resultValues
+
+                reducedValues.Free processor
+
+                let resultFirstKeys =
+                    clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                scatterIndices processor resultPositions firstReducedKeys resultFirstKeys
+
+                firstReducedKeys.Free processor
+
+                let resultSecondKeys =
+                    clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                scatterIndices processor resultPositions secondReducedKeys resultSecondKeys
+
+                secondReducedKeys.Free processor
+
+                resultPositions.Free processor
+
+                resultFirstKeys, resultSecondKeys, resultValues

tests/GraphBLAS-sharp.Tests/Common/Reduce/ReduceByKey.fs

@@ -2,6 +2,7 @@ module GraphBLAS.FSharp.Tests.Backend.Common.Reduce.ByKey
 
 open Expecto
 open GraphBLAS.FSharp.Backend.Common
+open GraphBLAS.FSharp.Backend.Quotes
 open GraphBLAS.FSharp.Test
 open GraphBLAS.FSharp.Tests
 open GraphBLAS.FSharp.Backend.Objects.ClContext
@@ -14,6 +15,16 @@ let processor = Context.defaultContext.Queue
 
 let config = Utils.defaultConfig
 
+let getOffsets array =
+    Array.map fst array
+    |> HostPrimitives.getUniqueBitmapFirstOccurrence
+    |> HostPrimitives.getBitPositions
+
+let getOffsets2D array =
+    Array.map (fun (fst, snd, _) -> fst, snd) array
+    |> HostPrimitives.getUniqueBitmapFirstOccurrence
+    |> HostPrimitives.getBitPositions
+
 let checkResult isEqual actualKeys actualValues keys values reduceOp =
 
     let expectedKeys, expectedValues =
@@ -336,3 +347,87 @@ let sequentialSegmentTests2D =
               createTestSequentialSegments2D<bool> (=) (&&) <@ (&&) @> ]
 
     testList "Sequential segments 2D" [ addTests; mulTests ]
+
+let checkResult2DOption isEqual firstActualKeys secondActualKeys actualValues firstKeys secondKeys values reduceOp =
+
+    let reduceOp left right =
+        match left, right with
+        | Some left, Some right ->
+            reduceOp left right
+        | Some value, None
+        | None, Some value -> Some value
+        | _ -> None
+
+    let expectedFirstKeys, expectedSecondKeys, expectedValues =
+        let keys = Array.zip firstKeys secondKeys
+
+        Array.zip keys values
+        |> Array.groupBy fst
+        |> Array.map (fun (key, array) -> key, Array.map snd array)
+        |> Array.map (fun (key, array) ->
+            Array.map Some array
+            |> Array.reduce reduceOp
+            |> fun result -> key, result)
+        |> Array.choose (fun ((fstKey, sndKey), value) ->
+            match value with
+            | Some value -> Some (fstKey, sndKey, value)
+            | _ -> None )
+        |> Array.unzip3
+
+    "First keys must be the same"
+    |> Utils.compareArrays (=) firstActualKeys expectedFirstKeys
+
+    "Second keys must be the same"
+    |> Utils.compareArrays (=) secondActualKeys expectedSecondKeys
+
+    "Values must the same"
+    |> Utils.compareArrays isEqual actualValues expectedValues
+
+let test2DOption<'a> isEqual reduce reduceOp (array: (int * int * 'a) [])  =
+    if array.Length > 0 then
+        let array = Array.sortBy (fun (fst, snd, _) -> fst, snd) array
+
+        let offsets = getOffsets2D array
+
+        let firstKeys, secondKeys, values = Array.unzip3 array
+
+        let clOffsets =
+            context.CreateClArrayWithSpecificAllocationMode(HostInterop, offsets)
+
+        let clFirstKeys =
+            context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, firstKeys)
+
+        let clSecondKeys =
+            context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, secondKeys)
+
+        let clValues =
+            context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, values)
+
+        let clFirstActualKeys, clSecondActualKeys, clReducedValues: ClArray<int> * ClArray<int> * ClArray<'a> =
+            reduce processor DeviceOnly offsets.Length clOffsets clFirstKeys clSecondKeys clValues
+
+        let reducedFirsKeys = clFirstActualKeys.ToHostAndFree processor
+        let reducesSecondKeys = clSecondActualKeys.ToHostAndFree processor
+        let reducedValues = clReducedValues.ToHostAndFree processor
+
+        checkResult2DOption isEqual reducedFirsKeys reducesSecondKeys reducedValues firstKeys secondKeys values reduceOp
+
+let createTest2DOption (isEqual: 'a -> 'a -> bool) (reduceOpQ, reduceOp) =
+    let reduce =
+        Reduce.ByKey2D.segmentSequentialOption context Utils.defaultWorkGroupSize reduceOpQ
+
+    test2DOption<'a> isEqual reduce reduceOp
+    |> testPropertyWithConfig { config with arbitrary = [ typeof<Generators.ArrayOfDistinctKeys> ] } $"test on {typeof<'a>}"
+
+let testsByKey2DSegmentsSequential =
+    [ createTest2DOption (=) ArithmeticOperations.intAdd
+
+      if Utils.isFloat64Available context.ClDevice then
+        createTest2DOption Utils.floatIsEqual ArithmeticOperations.floatAdd
+
+      createTest2DOption Utils.float32IsEqual ArithmeticOperations.float32Add
+      createTest2DOption (=) ArithmeticOperations.boolAdd ]
+    |> testList "2D option"
+
+
+

tests/GraphBLAS-sharp.Tests/Program.fs

@@ -95,10 +95,10 @@ let allTests =
     testList
         "All tests"
         [ // SpGeMM.expandTests
-          SpGeMM.generalTests
+          // SpGeMM.generalTests
           // Common.Gather.initTests
           // Common.ClArray.Choose.tests2 ]
-        ]
+          Common.Reduce.ByKey.testsByKey2DSegmentsSequential ]
     |> testSequenced
 
 [<EntryPoint>]