add: Search module

Author: IgorErin

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

@@ -5,6 +5,8 @@ open GraphBLAS.FSharp.Backend.Quotes
 open Microsoft.FSharp.Control
 open Microsoft.FSharp.Quotations
 open GraphBLAS.FSharp.Backend.Objects.ClContext
+open GraphBLAS.FSharp.Backend.Objects.ClCell
+open GraphBLAS.FSharp.Backend.Objects.ArraysExtensions
 
 module Reduce =
     let private runGeneral (clContext: ClContext) workGroupSize scan scanToCell =
@@ -235,3 +237,166 @@ module Reduce =
             runGeneral clContext workGroupSize scan scanToCell
 
         fun (processor: MailboxProcessor<_>) (array: ClArray<'a>) -> run processor array
+
+    module ByKey =
+        let sequential (clContext: ClContext) workGroupSize (reduceOp: Expr<'a -> 'a -> 'a>) =
+
+            let kernel =
+                <@ fun (ndRange: Range1D) length (keys: ClArray<int>) (values: ClArray<'a>) (reducedValues: ClArray<'a>) (reducedKeys: ClArray<int>) ->
+
+                    let gid = ndRange.GlobalID0
+
+                    if gid = 0  then
+                         let mutable currentKey = keys.[gid]
+                         let mutable segmentResult = values.[gid]
+                         let mutable segmentCount = 0
+
+                         for i in 1 .. length - 1 do
+                             if currentKey = keys.[i] then
+                                 segmentResult <- (%reduceOp) segmentResult values.[i]
+                             else
+                                 reducedValues.[segmentCount] <- segmentResult
+                                 reducedKeys.[segmentCount] <- currentKey
+
+                                 segmentCount <- segmentCount + 1
+                                 currentKey <- keys.[i]
+                                 segmentResult <- values.[i]
+
+                         reducedKeys.[segmentCount] <- currentKey
+                         reducedValues.[segmentCount] <- segmentResult @>
+
+            let kernel = clContext.Compile kernel
+
+            fun (processor: MailboxProcessor<_>) allocationMode (resultLength: int) (keys: ClArray<int>) (values: ClArray<'a>) ->
+
+                let reducedValues = clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let reducedKeys = clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let ndRange = Range1D.CreateValid(resultLength, workGroupSize)
+
+                let kernel = kernel.GetKernel()
+
+                processor.Post(Msg.MsgSetArguments(fun () -> kernel.KernelFunc ndRange resultLength keys values reducedValues reducedKeys))
+
+                processor.Post(Msg.CreateRunMsg<_, _>(kernel))
+
+        let segmentSequential (clContext: ClContext) workGroupSize (reduceOp: Expr<'a -> 'a -> 'a>) =
+
+            let kernel =
+                <@ fun (ndRange: Range1D) uniqueKeyCount (offsets: ClArray<int>) (keys: ClArray<int>) (values: ClArray<'a>) (reducedValues: ClArray<'a>) (reducedKeys: ClArray<int>) ->
+
+                    let gid = ndRange.GlobalID0
+
+                    if gid < uniqueKeyCount then
+                        let startPosition = offsets.[gid]
+                        let sourceKey = keys.[startPosition]
+
+                        let mutable nextPosition = startPosition + 1 // TODO()
+                        let mutable nextKey = keys.[nextPosition]
+                        let mutable sum = values.[startPosition]
+
+                        while nextKey = sourceKey do
+                            sum <- (%reduceOp) sum values.[nextPosition]
+
+                            nextPosition <- nextPosition + 1
+                            nextKey <- keys.[nextPosition]
+
+                        reducedValues.[gid] <- sum
+                        reducedKeys.[gid] <- sourceKey @>
+
+            let kernel = clContext.Compile kernel
+
+            let getUniqueBitmap = ClArray.getUniqueBitmap clContext workGroupSize
+
+            let prefixSum = PrefixSum.runExcludeInplace <@ (+) @> clContext workGroupSize
+
+            let removeDuplicates = ClArray.removeDuplications clContext workGroupSize
+
+            fun (processor: MailboxProcessor<_>) allocationMode (keys: ClArray<int>) (values: ClArray<'a>) ->
+
+                let bitmap = getUniqueBitmap processor DeviceOnly keys
+
+                let resultLength = (prefixSum processor bitmap 0).ToHostAndFree processor
+
+                let offsets = removeDuplicates processor bitmap
+
+                bitmap.Free processor
+
+                let reducedValues = clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let reducedKeys = clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let ndRange = Range1D.CreateValid(resultLength, workGroupSize)
+
+                let kernel = kernel.GetKernel()
+
+                processor.Post(Msg.MsgSetArguments(fun () -> kernel.KernelFunc ndRange resultLength offsets keys values reducedValues reducedKeys))
+
+                processor.Post(Msg.CreateRunMsg<_, _>(kernel))
+
+        let oneWorkGroupSegments (clContext: ClContext) workGroupSize (reduceOp: Expr<'a -> 'a -> 'a>) =
+
+            let kernel =
+                <@ fun (ndRange: Range1D) length (keys: ClArray<int>) (values: ClArray<'a>) (reducedValues: ClArray<'a>) (reducedKeys: ClArray<int>) ->
+
+                    let lid = ndRange.GlobalID0
+
+                    // load values to local memory (may be without it)
+                    let localValues = localArray<'a> length
+                    if lid < length then localValues.[lid] <- values.[lid]
+
+                    // load keys to local memory (mb without it)
+                    let localKeys = localArray<int> length
+                    if lid < length then localKeys.[lid] <- keys.[lid]
+
+                    // get unique keys bitmap
+                    let localBitmap = localArray<int> length
+                    (%PreparePositions.getUniqueBitmapLocal<int>) localKeys length lid localBitmap
+
+                    // get positions from bitmap by prefix sum
+                    // ??? get bitmap by prefix sum in another kernel ???
+                    (%SubSum.localIntPrefixSum) lid workGroupSize localBitmap
+                    let localPositions = localBitmap
+
+                    let uniqueKeysCount = localPositions.[length - 1]
+
+                    if lid < uniqueKeysCount then
+                        let itemKeyId = lid + 1
+                        // we can count start position by itemKeyId
+                        // but loose coalesced memory read pattern
+
+                        let startKeyIndex =
+                            (%Search.Bin.lowerPosition) length itemKeyId localPositions
+
+                        match startKeyIndex with
+                        | Some startPosition ->
+                            let sourcePosition = localPositions.[startPosition]
+                            let mutable currentSum = localValues.[startPosition]
+                            let mutable currentIndex = startPosition + 1
+
+                            while currentIndex < length
+                                  && localPositions.[currentIndex]  = sourcePosition do
+
+                                currentSum <- (%reduceOp) currentSum localValues.[currentIndex]
+                                currentIndex <- currentIndex + 1
+
+                            reducedKeys.[lid] <- localKeys.[startPosition]
+                            reducedValues.[lid] <- currentSum
+                        | None -> () @>
+
+            let kernel = clContext.Compile kernel
+
+            fun (processor: MailboxProcessor<_>) allocationMode (resultLength: int) (keys: ClArray<int>) (values: ClArray<'a>) ->
+
+                let reducedValues = clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let reducedKeys = clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
+
+                let ndRange = Range1D.CreateValid(resultLength, workGroupSize)
+
+                let kernel = kernel.GetKernel()
+
+                processor.Post(Msg.MsgSetArguments(fun () -> kernel.KernelFunc ndRange resultLength keys values reducedValues reducedKeys))
+
+                processor.Post(Msg.CreateRunMsg<_, _>(kernel))

src/GraphBLAS-sharp.Backend/GraphBLAS-sharp.Backend.fsproj

@@ -26,11 +26,12 @@
     <Compile Include="Quotes/PreparePositions.fs" />
     <Compile Include="Quotes/Predicates.fs" />
     <Compile Include="Quotes/Map.fs" />
+    <Compile Include="Quotes\Search.fs" />
     <Compile Include="Common/Scatter.fs" />
     <Compile Include="Common/Utils.fs" />
-    <Compile Include="Common/Sum.fs" />
     <Compile Include="Common/PrefixSum.fs" />
     <Compile Include="Common/ClArray.fs" />
+    <Compile Include="Common\Sum.fs" />
     <Compile Include="Common/BitonicSort.fs" />
     <Compile Include="Predefined/PrefixSum.fs" />
     <!--Compile Include="Matrices.fs" /-->

src/GraphBLAS-sharp.Backend/Matrix/COOMatrix/Map2.fs

@@ -7,33 +7,9 @@ open GraphBLAS.FSharp.Backend.Objects
 open GraphBLAS.FSharp.Backend
 open GraphBLAS.FSharp.Backend.Objects.ClMatrix
 open GraphBLAS.FSharp.Backend.Objects.ClContext
+open GraphBLAS.FSharp.Backend.Quotes
 
 module internal Map2 =
-    let binSearch<'a> =
-        <@ fun lenght sourceIndex (rowIndices: ClArray<int>) (columnIndices: ClArray<int>) (values: ClArray<'a>) ->
-
-            let mutable leftEdge = 0
-            let mutable rightEdge = lenght - 1
-
-            let mutable result = None
-
-            while leftEdge <= rightEdge do
-                let middleIdx = (leftEdge + rightEdge) / 2
-
-                let currentIndex: uint64 =
-                    ((uint64 rowIndices.[middleIdx]) <<< 32)
-                    ||| (uint64 columnIndices.[middleIdx])
-
-                if sourceIndex = currentIndex then
-                    result <- Some values.[middleIdx]
-
-                    rightEdge <- -1 // TODO() break
-                elif sourceIndex < currentIndex then
-                    rightEdge <- middleIdx - 1
-                else
-                    leftEdge <- middleIdx + 1
-
-            result @>
 
     let preparePositions<'a, 'b, 'c> (clContext: ClContext) workGroupSize opAdd =
 
@@ -51,10 +27,10 @@ module internal Map2 =
                         (uint64 rowIndex <<< 32) ||| (uint64 columnIndex)
 
                     let leftValue =
-                        (%binSearch) leftValuesLength index leftRows leftColumns leftValues
+                        (%Search.Bin.byKey2) leftValuesLength index leftRows leftColumns leftValues
 
                     let rightValue =
-                        (%binSearch) rightValuesLength index rightRows rightColumn rightValues
+                        (%Search.Bin.byKey2) rightValuesLength index rightRows rightColumn rightValues
 
                     match (%op) leftValue rightValue with
                     | Some value ->

src/GraphBLAS-sharp.Backend/Objects/ArraysExtentions.fs

@@ -13,6 +13,14 @@ module ArraysExtensions =
             let dst = Array.zeroCreate this.Length
             q.PostAndReply(fun ch -> Msg.CreateToHostMsg(this, dst, ch))
 
+        member this.Free(q: MailboxProcessor<_>) = q.Post <| Msg.CreateFreeMsg this
+
+        member this.ToHostAndFree(q: MailboxProcessor<_>) =
+            let result = this.ToHost q
+            this.Free q
+
+            result
+
         member this.Size = this.Length
 
     type 'a ``[]`` with

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

@@ -27,3 +27,12 @@ module PreparePositions =
                     allValuesBuffer.[index] <- v
                     rawPositionsBuffer.[index] <- 1
                 | None -> rawPositionsBuffer.[index] <- 0 @>
+
+    let getUniqueBitmapLocal<'a when 'a : equality> =
+        <@ fun (array: 'a []) length lid (result: int []) ->
+            if lid < length then
+                let isFirst = lid = 0
+                let isUnique = lid > 0 && array.[lid] <> array.[lid - 1]
+
+                if isFirst || isUnique then result.[lid] <- 1 else result.[lid] <- 0 @>
+

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

@@ -0,0 +1,80 @@
+namespace GraphBLAS.FSharp.Backend.Quotes
+
+open Brahma.FSharp
+
+module Search =
+    module Bin =
+        let byKey<'a> =
+            <@ fun lenght sourceIndex (indices: ClArray<int>) (values: ClArray<'a>) ->
+
+                let mutable leftEdge = 0
+                let mutable rightEdge = lenght - 1
+
+                let mutable result = None
+
+                while leftEdge <= rightEdge do
+                    let middleIdx = (leftEdge + rightEdge) / 2
+                    let currentIndex = indices.[middleIdx]
+
+                    if sourceIndex = currentIndex then
+                        result <- Some values.[middleIdx]
+
+                        rightEdge <- -1 // TODO() break
+                    elif sourceIndex < currentIndex then
+                        rightEdge <- middleIdx - 1
+                    else
+                        leftEdge <- middleIdx + 1
+
+                result @>
+
+        let byKey2<'a> =
+            <@ fun lenght sourceIndex (rowIndices: ClArray<int>) (columnIndices: ClArray<int>) (values: ClArray<'a>) ->
+
+                let mutable leftEdge = 0
+                let mutable rightEdge = lenght - 1
+
+                let mutable result = None
+
+                while leftEdge <= rightEdge do
+                    let middleIdx = (leftEdge + rightEdge) / 2
+
+                    let currentIndex: uint64 =
+                        ((uint64 rowIndices.[middleIdx]) <<< 32)
+                        ||| (uint64 columnIndices.[middleIdx])
+
+                    if sourceIndex = currentIndex then
+                        result <- Some values.[middleIdx]
+
+                        rightEdge <- -1 // TODO() break
+                    elif sourceIndex < currentIndex then
+                        rightEdge <- middleIdx - 1
+                    else
+                        leftEdge <- middleIdx + 1
+
+                result @>
+
+        /// <summary>
+        /// Find lower position of item in array.
+        /// </summary>
+        let lowerPosition<'a when 'a : equality and 'a: comparison> =
+            <@ fun lenght sourceItem (keys: '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 @>

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

@@ -4,7 +4,7 @@ open Brahma.FSharp
 
 module SubSum =
     let private treeAccess<'a> opAdd =
-        <@ fun step lid wgSize (localBuffer: 'a []) ->
+        <@ fun step lid _ (localBuffer: 'a []) ->
             let i = step * (lid + 1) - 1
 
             let firstValue = localBuffer.[i - (step >>> 1)]
@@ -35,3 +35,21 @@ module SubSum =
         sumGeneral<'a> <| sequentialAccess<'a> opAdd
 
     let treeSum<'a> opAdd = sumGeneral<'a> <| treeAccess<'a> opAdd
+
+    let localPrefixSum opAdd =
+        <@ fun (lid: int) (workGroupSize: int) (array: 'a []) ->
+            let mutable offset = 1
+
+            while offset < workGroupSize do
+                barrierLocal ()
+                let mutable value = array.[lid]
+
+                if lid >= offset then
+                    value <- (%opAdd) value array.[lid - offset]
+
+                offset <- offset * 2
+
+                barrierLocal ()
+                array.[lid] <- value @>
+
+    let localIntPrefixSum = localPrefixSum <@ (+) @>