SegReduce without offsets

Author: kirillgarbar

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

@@ -428,7 +428,7 @@ module Reduce =
                         let itemKeyId = lid + 1
 
                         let startKeyIndex =
-                            (%Search.Bin.lowerPosition) length itemKeyId localBitmap
+                            (%Search.Bin.lowerPositionLocal) length itemKeyId localBitmap
 
                         match startKeyIndex with
                         | Some startPosition ->
@@ -473,16 +473,144 @@ module Reduce =
                 reducedValues, reducedKeys
 
         module Option =
+            /// <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>
+            let segmentSequential<'a>
+                (reduceOp: Expr<'a option -> 'a option -> 'a option>)
+                (clContext: ClContext)
+                workGroupSize
+                =
+
+                let kernel =
+                    <@ fun (ndRange: Range1D) uniqueKeyCount keysLength (offsets: ClArray<int>) (keys: ClArray<int>) (values: ClArray<'a option>) (reducedValues: ClArray<'a>) (firstReducedKeys: ClArray<int>) (resultPositions: ClArray<int>) ->
+
+                        let gid = ndRange.GlobalID0
+
+                        if gid < uniqueKeyCount then
+                            let startPosition =
+                                (%Search.Bin.lowerPosition) keysLength gid offsets
+
+                            match startPosition with
+                            | Some startPosition ->
+                                let firstSourceKey = keys.[startPosition]
+
+                                let mutable sum = None
+
+                                let mutable currentPosition = startPosition
+
+                                while currentPosition < keysLength
+                                      && firstSourceKey = keys.[currentPosition] do
+                                    let result = (%reduceOp) sum values.[currentPosition] // brahma error
+                                    sum <- result
+                                    currentPosition <- currentPosition + 1
+
+                                match sum with
+                                | Some value ->
+                                    reducedValues.[gid] <- value
+                                    resultPositions.[gid] <- 1
+                                | None -> resultPositions.[gid] <- 0
+
+                                firstReducedKeys.[gid] <- firstSourceKey
+                            | None -> () @> // not possible if done correctly
+
+                let kernel = clContext.Compile kernel
+
+                let getUniqueBitmap =
+                    ClArray.Bitmap.lastOccurrence clContext workGroupSize
+
+                let scatterData =
+                    Scatter.lastOccurrence clContext workGroupSize
+
+                let scatterIndices =
+                    Scatter.lastOccurrence clContext workGroupSize
+
+                let prefixSum =
+                    PrefixSum.standardExcludeInPlace clContext workGroupSize
+
+                fun (processor: MailboxProcessor<_>) allocationMode (keys: ClArray<int>) (values: ClArray<'a option>) ->
+
+                    let offsets =
+                        getUniqueBitmap processor DeviceOnly keys
+
+                    let uniqueKeysCount =
+                        (prefixSum processor offsets)
+                            .ToHostAndFree processor
+
+                    let reducedValues =
+                        clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, uniqueKeysCount)
+
+                    let reducedKeys =
+                        clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, uniqueKeysCount)
+
+                    let resultPositions =
+                        clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, uniqueKeysCount)
+
+                    let ndRange =
+                        Range1D.CreateValid(uniqueKeysCount, workGroupSize)
+
+                    let kernel = kernel.GetKernel()
+
+                    processor.Post(
+                        Msg.MsgSetArguments
+                            (fun () ->
+                                kernel.KernelFunc
+                                    ndRange
+                                    uniqueKeysCount
+                                    keys.Length
+                                    offsets
+                                    keys
+                                    values
+                                    reducedValues
+                                    reducedKeys
+                                    resultPositions)
+                    )
+
+                    processor.Post(Msg.CreateRunMsg<_, _>(kernel))
+
+                    let resultLength =
+                        (prefixSum processor resultPositions)
+                            .ToHostAndFree processor
+
+                    if resultLength = 0 then
+                        None
+                    else
+                        // write values
+                        let resultValues =
+                            clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                        scatterData processor resultPositions reducedValues resultValues
+
+                        reducedValues.Free processor
+
+                        // write keys
+                        let resultKeys =
+                            clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                        scatterIndices processor resultPositions reducedKeys resultKeys
+
+                        reducedKeys.Free processor
+                        resultPositions.Free processor
+
+                        Some(resultValues, resultKeys)
+
             /// <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.
+            /// The length of the result and offsets for each segment must be calculated in advance.
             /// </remarks>
-            let segmentSequential<'a> (reduceOp: Expr<'a -> 'a -> 'a option>) (clContext: ClContext) workGroupSize =
+            let segmentSequentialByOffsets<'a>
+                (reduceOp: Expr<'a -> 'a -> 'a option>)
+                (clContext: ClContext)
+                workGroupSize
+                =
 
                 let kernel =
                     <@ fun (ndRange: Range1D) uniqueKeyCount keysLength (offsets: ClArray<int>) (keys: ClArray<int>) (values: ClArray<'a>) (reducedValues: ClArray<'a>) (firstReducedKeys: ClArray<int>) (resultPositions: ClArray<int>) ->

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

@@ -98,6 +98,32 @@ module Search =
         /// Find lower position of item in array.
         /// </summary>
         let lowerPosition<'a when 'a: equality and 'a: comparison> =
+            <@ fun lenght sourceItem (keys: ClArray<'a>) ->
+
+                let mutable leftEdge = 0
+                let mutable rightEdge = lenght - 1
+                let mutable resultPosition = None
+
+                while leftEdge <= rightEdge do
+                    let currentPosition = (leftEdge + rightEdge) / 2
+                    let currentKey = keys.[currentPosition]
+
+                    if sourceItem = currentKey then
+                        // remember positions and move left
+                        resultPosition <- Some currentPosition
+
+                        rightEdge <- currentPosition - 1
+                    elif sourceItem < currentKey then
+                        rightEdge <- currentPosition - 1
+                    else
+                        leftEdge <- currentPosition + 1
+
+                resultPosition @>
+
+        /// <summary>
+        /// Find lower position of item in array.
+        /// </summary>
+        let lowerPositionLocal<'a when 'a: equality and 'a: comparison> =
             <@ fun lenght sourceItem (keys: 'a []) ->
 
                 let mutable leftEdge = 0

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

@@ -430,7 +430,7 @@ let testOption<'a> isEqual reduceOp testFun (array: (int * 'a) []) =
         |> checkResultOption isEqual keys values reduceOp
 
 let createTestOption (isEqual: 'a -> 'a -> bool) (reduceOpQ, reduceOp) =
-    Reduce.ByKey.Option.segmentSequential reduceOpQ context Utils.defaultWorkGroupSize
+    Reduce.ByKey.Option.segmentSequentialByOffsets reduceOpQ context Utils.defaultWorkGroupSize
     |> testOption<'a> isEqual reduceOp
     |> testPropertyWithConfig
         { config with